BEVFormer是一种纯视觉的自动驾驶感知算法,通过融合环视相机图像的空间和时序特征显式的生成具有强表征能力的BEV特征,并应用于下游3D检测、分割等任务,取得了SOTA的结果。我们在EasyCV开源框架(https://github.com/alibaba/Ea…)中,对BEVFomer算法进行集成,并从训练速度、算法收敛速度角度对代码进行了一些优化。同时,我们进一步使用推理优化工具PAI-Blade对模型进行优化,相比于原始模型在A100配置下能取得40%的推理速度提升。本文将从以下几个部分进行介绍:1、BEVFormer算法思想 2、训练速度和算法收敛速度优化 3、使用PAI-Blade优化推理速度。
BEVFormer 算法思想
如上图所示,BEVFormer由如下三个部分组成:
backbone:用于从6个角度的环视图像中提取多尺度的multi-camera feature
BEV encoder:该模块主要包括Temporal self-Attention 和 Spatial Cross-Attention两个部分。
Spatial Cross-Attention结合多个相机的内外参信息对对应位置的multi-camera feature进行query,从而在统一的BEV视角下将multi-camera feature进行融合。
Temporal self-Attention将History BEV feature和 current BEV feature
通过 self-attention module进行融合。通过上述两个模块,输出同时包含多视角和时序信息的BEV feature进一步用于下游3D检测和分割任务
Det&Seg Head:用于特定任务的task head
BEVFormer训练优化
训练加速优化
我们从数据读取和减少内存拷贝消耗等角度对训练代码进行优化。
数据读取
使用更高效的图片解码库 turbojpeg
BEVFormer在训练过程中,需要时序上的数据作为输入,将串形的读取方式优化为并行读取。
先做resize再做其他预处理,减少了额外像素带来的计算开销
内存拷贝优化
使用pin_memery=True,并修复了mmcv DataContainer pin_memory的bug
将代码中的numpy操作替换为torch.tensor,避免不必要的h2d拷贝
other
使用torch.backends.cudnn.benchmark=True(ps:需要保证在输入数据没有动态性的情况下使用,否则反而会增加训练耗时)
修复了torch.cuda.amp混合精度在LayerNorm层失效的bug
我们在A100 80G的机器上,使用fp16对比吞吐量如下:
精度收敛优化
我们使用额外的数据增广方式和不同的损失函数来优化模型。同时加入额外的训练策略来进一步提升模型收敛速度及精度。
数据增广方式
rand scale(采用不同分辨率的输入进行训练,实验中发现该操作会引入至少20%的额外训练时间,因此在下述实验中,均没有采用)
rand_flip(以50%的概率随机翻转图片)
损失函数
使用smooth l1 loss或 balance l1 loss代替l1 loss。(在mini dataset的实验中,这两个损失都可以提升精度,下面的实验中采用balance l1 loss)
训练策略
1.使用one2many Branch
这个做法来自于H-Deformable-DETR,在DETR系列的检测模型中采用one2one的匹配方式来分配GT Boxes,这种做法虽然让模型在测试的时候,能够避免冗余的NMS后处理操作,但是只有少数的Query会被分配给正样本,导致训练时模型收敛速度相比于one2many的方式会慢很多。因此,在训练过程中加入auxiliary Query,同一个GT Box会匹配多个auxiliary Query,并使用attention mask将one2one branch和one2many branch的信息隔离开。通过这样的方式,能够显著的提升训练过程中的收敛速度,同时在测试过程中只需要保持one2one branch进行预测。(在实验中,使用额外加入1800个auxiliary Query,每个GT box匹配4个query进行训练)
CBGS in one2many Branch
我们的实验是在NuScenes数据集上进行的,在该数据集的3D检测任务上有10类标签,但是这10类标签之间的样本极度不均衡,很多算法会采用CBGS操作进行类间样本均衡,但是这个操作会将整个数据集扩大4.5倍,虽然有一定的精度提升,但是也带来了巨大的训练成本。我们考虑在one2many Branch上进行样本均衡操作,即对于实例数量较多的样本使用较少的auxiliary Query进行匹配,而对于长尾的样本使用较多的auxiliary Query进行匹配。通过CBGS in one2many Branch的方式,训练时间和base保持一致的基础上会进一步提升收敛速度,最终的精度也有一定的提升。(实验中匹配框数量变化:[4, 4, 4, 4, 4, 4, 4, 4, 4, 4] -> [2, 3, 7, 7, 9, 6, 7, 6, 2, 5])
我们在单机8卡A100 80G下进行实验,如下表所示:
模型收敛速度如下图所示:
由上图可以看出,使用上述优化方式可以大幅提升模型收敛速度,仅需要75%的训练时间就可以达到base的最终精度。同时最终的NDS相比于base也有0.8的提升。
详细配置,训练log和模型权重,参考:
https://www.xiaohongshu.com/discovery/item/63189ac1000000001103426a
https://www.xiaohongshu.com/discovery/item/63161bdf0000000011013a26
https://www.xiaohongshu.com/discovery/item/6314c2ad0000000008021a71
https://www.xiaohongshu.com/discovery/item/6314c07a00000000120098b5
https://www.xiaohongshu.com/discovery/item/6314bf88000000000900c7ef
https://www.xiaohongshu.com/discovery/item/6313649d000000001101079e
在阿里云机器学习平台PAI上使用BEVFormer模型
PAI-DSW(Data Science Workshop)是阿里云机器学习平台PAI开发的云上IDE,面向各类开发者,提供了交互式的编程环境。在DSW Gallery中(链接),提供了各种Notebook示例,方便用户轻松上手DSW,搭建各种机器学习应用。我们也在DSW Gallery中上架了BEVFormer进行3D检测的Sample Notebook(见下图),欢迎大家体验!
使用PAI-Blade进行推理加速
PAI-Blade是由阿里云机器学习平台PAI开发的模型优化工具,可以针对不同的设备不同模型进行推理加速优化。PAI-Blade遵循易用性,鲁棒性和高性能为原则,将模型的部署优化进行高度封装,设计了统一简单的API,在完成Blade环境安装后,用户可以在不了解ONNX、TensorRT、编译优化等技术细节的条件下,通过简单的代码调用方便的实现对模型的高性能部署。更多PAI-Blade相关技术介绍可以参考 [PAI-Blade介绍]。
PAI-EasyCV中对Blade进行了支持,用户可以通过PAI-EasyCV的训练config 中配置相关export 参数,从而对训练得到的模型进行导出。
对于BEVFormer模型,我们在A100机器下进行进行推理速度对比,使用PAI-Blade优化后的模型能取得42%的优化加速。
环境准备
我们提供一个PAI-Blade + PAI-EasyCV 的镜像包供用户可以直接使用,镜像包地址:easycv-blade-torch181-cuda111.tar
用户也可以基于Blade每日发布的镜像自行搭建推理环境 [PAI-Blade社区镜像发布]。
自行搭建环境时需要注意:BEVFomer-base使用resnet101-dcn作为image backbone,DCN算子使用的是mmcv中的自定义算子,为了导出TorchScript,我们对该接口进行了修改。所以mmcv需要源码编译。