基于 EasyCV 加速 BEVFormer

申屠洛华

2023-12-01

导言

BEVFormer 是一种纯视觉的自动驾驶感知算法，通过融合环视相机图像的空间和时序特征显式的生成具有强表征能力的 BEV 特征，并应用于下游 3D 检测、分割等任务，取得了 SOTA 的结果。我们在 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 对比吞吐量如下：

Setting	throughput(samples/s)
BEVFormer-tiny bs=32	3.55
EasyCV BEVFormer-tiny bs=32	9.84( +177% )
BEVFormer-base bs=5	0.727
EasyCV BEVFormer-base bs=5	0.8( +10% )

精度收敛优化

我们使用额外的数据增广方式和不同的损失函数来优化模型。同时加入额外的训练策略来进一步提升模型收敛速度及精度。

数据增广方式
- 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 下进行实验，如下表所示：

config setting	NDS	mAP	throughput(samples/s)
官方 BEVFormer-base	52.44	41.91	3.289
EasyCV BEVFormer-base	52.66	42.13	3.45
EasyCV BEVFormer-base-one2manybranch	53.02(+0.58)	42.48(+0.57)	3.40
EasyCV BEVFormer-base-cbgs_one2manybranch	53.28(+0.84)	42.63(+0.72)	3.41

模型收敛速度如下图所示：

由上图可以看出，使用上述优化方式可以大幅提升模型收敛速度，仅需要 75% 的训练时间就可以达到 base 的最终精度。同时最终的 NDS 相比于 base 也有 0.8 的提升。

在阿里云机器学习平台 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% 的优化加速。

Name	Backend	Median(FPS)	Mean(FPS)	Median(ms)	Mean(ms)
easycv	TensorRT	3.68697	3.68651	0.271226	0.271259
easycv script	TensorRT	3.8131	3.79859	0.262254	0.26337
blade	TensorRT	5.40248	5.23383(+42%)	0.1851	0.192212

环境准备

我们提供一个 PAI-Blade + PAI-EasyCV 的镜像包供用户可以直接使用，镜像包地址：easycv-blade-torch181-cuda111.tar

用户也可以基于 Blade 每日发布的镜像自行搭建推理环境 [PAI-Blade 社区镜像发布]。

自行搭建环境时需要注意：BEVFomer-base 使用 resnet101-dcn 作为 image backbone，DCN 算子使用的是 mmcv 中的自定义算子，为了导出 TorchScript，我们对该接口进行了修改。所以 mmcv 需要源码编译。

clone mmcv 源码

$ git clone https://github.com/open-mmlab/mmcv.git

替换 mmcv 文件

替换时请注意 mmcv 的版本，注意接口要匹配。mmcv1.6.0 版本已验证。

源码编译

导出 Blade 模型

导出 Blade 的模型的配置可以参考文件 bevformer_base_r101_dcn_nuscenes.py 中的 export 字段，配置如下：

export = dict(
    type='blade',
    blade_config=dict(
        enable_fp16=True,
        fp16_fallback_op_ratio=0.0,
        customize_op_black_list=[
            'aten::select', 'aten::index', 'aten::slice', 'aten::view',
            'aten::upsample', 'aten::clamp'
        ]
    )
)

导出命令：

$ cd ${EASYCV_ROOT}
$ export PYTHONPATH='./'
$ python tools/export.py configs/detection3d/bevformer/bevformer_base_r101_dcn_nuscenes.py bevformer_base.pth bevformer_export.pth

Blade 模型推理

推理脚本：

from easycv.predictors import BEVFormerPredictor

blade_model_path = 'bevformer_export.pth.blade'
config_file = 'configs/detection3d/bevformer/bevformer_base_r101_dcn_nuscenes.py'

predictor = BEVFormerPredictor(
            model_path=blade_model_path,
            config_file=config_file,
            model_type='blade',
			)

inputs_file = 'nuscenes_infos_temporal_val.pkl'  # 以NuScenes val数据集文件为例
input_samples = mmcv.load(inputs_file)['infos']
predict_results = predictor(input_samples)
print(predict_results)

展望

我们在 EasyCV 框架中，集成了 BEVFormer 算法，并从训练加速、精度收敛和推理加速角度对算法进行了一些改进。近期，也涌现了许多新的 BEV 感知算法，如 BEVFormerv2。在 BEVFormerv2 中通过 Perspective Supervision 的方式，让算法能够不受限于使用一些在深度估计或 3D 检测上的预训练 backbone，而直接使用近期更有效的大模型 BackBone（如 ConvNext、DCNv3 等），同时采用 two-stage 的检测方式进一步增强模型能力，在 Nuscenes 数据集的 camera-based 3D 检测任务取得 sota 的结果。