BEVFormer-accelerate:基于 EasyCV 加速 BEVFormer
作者:贺弘 夕陌 谦言 临在
导言
BEVFormer 是一种纯视觉的自动驾驶感知算法,通过融合环视相机图像的空间和时序特征显式的生成具有强表征能力的 BEV 特征,并应用于下游 3D 检测、分割等任务,取得了 SOTA 的结果。我们在 EasyCV 开源框架(https://github.com/alibaba/EasyCV)中,对 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)
训练策略
使用 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://github.com/alibaba/EasyCV/blob/master/docs/source/model_zoo_det3d.md
在阿里云机器学习平台 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 需要源码编译。
clone mmcv 源码
替换 mmcv 文件
替换时请注意 mmcv 的版本,注意接口要匹配。mmcv1.6.0 版本已验证。
参考easycv/thirdparty/mmcv/目录下的修改文件。用mmcv/ops/csrc/pytorch/modulated_deform_conv.cpp和mmcv/ops/modulated_deform_conv.py去替换 mmcv 中的原文件。
源码编译
mmcv 源码编译请参考:https://mmcv.readthedocs.io/en/latest/get_started/build.html
导出 Blade 模型
导出 Blade 的模型的配置可以参考文件bevformer_base_r101_dcn_nuscenes.py中的 export 字段,配置如下:
导出命令:
Blade 模型推理
推理脚本:
NuScenes 数据集准备请参考:NuScenes数据集准备
展望
我们在 EasyCV 框架中,集成了 BEVFormer 算法,并从训练加速、精度收敛和推理加速角度对算法进行了一些改进。近期,也涌现了许多新的 BEV 感知算法,如 BEVFormerv2。在 BEVFormerv2 中通过 Perspective Supervision 的方式,让算法能够不受限于使用一些在深度估计或 3D 检测上的预训练 backbone,而直接使用近期更有效的大模型 BackBone(如 ConvNext、DCNv3 等),同时采用 two-stage 的检测方式进一步增强模型能力,在 Nuscenes 数据集的 camera-based 3D 检测任务取得 sota 的结果。
EasyCV(https://github.com/alibaba/EasyCV)会持续跟进业界 sota 方法,欢迎大家关注和使用,欢迎大家各种维度的反馈和改进建议以及技术讨论,同时我们十分欢迎和期待对开源社区建设感兴趣的同行一起参与共建。
EasyCV 往期分享
EasyCV 开源地址:https://github.com/alibaba/EasyCV
使用 EasyCV Mask2Former 轻松实现图像分割 https://zhuanlan.zhihu.com/p/585259059
EasyCV DataHub 提供多领域视觉数据集下载,助力模型生产 https://zhuanlan.zhihu.com/p/572593950
EasyCV 带你复现更好更快的自监督算法-FastConvMAE https://zhuanlan.zhihu.com/p/566988235
基于 EasyCV 复现 DETR 和 DAB-DETR,Object Query 的正确打开方式 https://zhuanlan.zhihu.com/p/543129581
基于 EasyCV 复现 ViTDet:单层特征超越 FPN https://zhuanlan.zhihu.com/p/528733299
MAE 自监督算法介绍和基于 EasyCV 的复现 https://zhuanlan.zhihu.com/p/515859470
EasyCV 开源|开箱即用的视觉自监督+Transformer 算法库 https://zhuanlan.zhihu.com/p/50521999
版权声明: 本文为 InfoQ 作者【阿里云大数据AI技术】的原创文章。
原文链接:【http://xie.infoq.cn/article/b91ab13ff60025f7897720463】。文章转载请联系作者。
评论