萤火跑模型 | 多卡并行助力 BEVFormer 三小时全精度训练

Freja    September 13, 2022

感知系统是自动驾驶系统中不可或缺的部分,系统中包括许多视觉任务,例如目标检测分割、目标跟踪、语义和实例分割、车道线检测等等。感知系统通过对周围环境状态的准确估计,为预测、规划等下游任务提供重要的依据。其中如何利用车载环视相机采集到的多个视角图像,实现精准的 3D 目标检测和分割是感知系统中最为基础且重要的任务。

2021年9月特斯拉的 AI Day 上,Karpathy 展示了他们的纯视觉自动驾驶感知框架,核心是基于 Transformer 网络结构,在鸟瞰图(BEV,Bird’s eye view)视角下进行感知,紧随其后,学术界也逐渐涌现了基于 Transformer 和 BEV 的优秀工作,其中 BEVFormer 通过构建时空融合网络,在 nuScenes 和 Waymo 的检测比赛上都榜上有名,受到了诸多关注。

然而,最近不少学者和开发者在复现优化 BEVFormer 模型时发现,当使用多显卡加速模型的训练时,随着并行 GPU 数的增加,训练最终的结果反而变差了。这无疑对自动驾驶模型的产业落地造成了不少制约。

幻方 AI 最近对 BEVFormer 进行了整理和优化,分析了其多显卡并行训练精度降低的原因。通过幻方自研的 3FSFFRecord 对模型训练进行提速;利用 hfai.autotune 进行自动集群参数的搜索,最终在幻方萤火集群上耗时差不多三小时完成了 BEVFormer 的大规模多卡并行训练,并实现了最终精度的复现

本期文章将为大家详细描述我们的训练优化过程,为研究者和开发者们提供效率提升的思路。

论文标题:BEVFormer: Learning Bird’s-Eye-View Representation from Multi-Camera Images via Spatiotemporal Transformers

原文地址:https://arxiv.org/abs/2203.17270

项目主页:https://github.com/HFAiLab/BEVFormer

BEVFormer 模型介绍

BEVFormer 模型基于可变形注意力机制,融合多视角相机到 BEV(Bird’s eye view)视角进行感知,适用于多种自动驾驶感知任务。该模型采用稀疏注意力机制提高了计算的效率,通过融合时空信息获得了更加准确的感知任务性能。BEVFormer 在 3 月份的 nuScenes 3D 目标检测任务上获得了第一名的成绩,mAP 和 NDS 两个指标均超过以往方法 3 个点以上。在 Waymo 组织的 Camera Track 比赛中,BEVFormer 模型同样获得了冠军。BEVFormer 模型整体架构如下图所示:

BEVFormer 框架图

BEVFormer 包括三个关键模块,第一个关键模块是 BEV 空间的查询向量(BEV Queries)。

模型设计了 BEV 视角下的网格特征,初始时为每个网格生成查询向量,通过对时序信息和空间信息的查询,将时空特征聚合到 BEV 查询向量的特征中,该特征为 3D 目标检测和地图语义分割的多种自动驾驶感知任务提供了丰富的感知依据。

相机和 BEV 视角对比图

为什么选择聚合到 BEV 视角?传感器的观测视角为相机视角,如上图左边图所示,自动驾驶汽车往往在车身装备多个摄像机,从而组成环视视角。BEV 视角相当于在车辆正上方 10-20 米的高度俯视进行观察,如上图右边图所示。

相比相机视角,使用 BEV 视角有两个好处,首先下游的规划、决策等任务都需要在车辆所处的 3D 世界下进行,而相机视角是 2D 像素世界,两者存在维度不匹配的问题,因此将多传感器数据融合到 BEV 视角,更加有利于下游模块的执行。其次,不同的物体在 BEV 视角下没有变形问题、视觉遮挡问题、物体重叠问题等,能够使得模型集中于解决感知问题。

第二个关键模块是空间交叉注意力模型(Spatial Cross-Attention)。

BEVFormer 通过 Transformer 结构融合多个视角的图片特征。在每个时间刻,对于每个栅格的 BEV 查询向量,将其所在位置投影到像素平面获得参考点,在参考点周围进行加权特征采样后更新 BEV 查询向量,从而完成空间的特征聚合。

该模块通过 Transformer 实现了从相机视角到 BEV 视角的特征图空间变化。特斯拉公布的方法是将每一个 BEV 查询向量,与各个视角下的所有特征计算注意力,密集的操作导致了过高的内存开销和计算复杂度。而 BEVFormer 通过将 BEV 视角位置投影到像素视角位置,进行对应位置的稀疏注意力计算,融合多视角的同时,大大降低了计算量和内存开销,与此同时,所学习到的注意力初始分布更加集中,梯度的有效性更高,有利于模型后续学习。

第三个关键模块是时序自注意力模型(Temporal Self-Attention)。

每一时刻的 BEV 特征都与上一时刻对齐后的 BEV 特征进行自注意力计算,以保证时序特征的动态更新。在自动驾驶感知任务中,时序信息能够弥补空间中发生的观测限制,并且提供感知所需要的运动信息。具体而言,首先时序信息是空间信息的很好的补充,当物体重叠导致暂时的遮挡、物体过大导致只能获得部分观测时,模型可以融合前后的时序信息获得更好的特征表示,以更加准确的完成感知任务。此外,时序信息可以提供物体的运动状态,例如物体的运动速度、运动方向等等,来帮助预测、规划等下游任务更好的进行。

时空融合示意图

如上图所示,最左边初始化了 BEV 视角下的栅格图,对于每一个位置,初始化一个 BEV 查询向量。BEV 查询向量一方面通过空间交叉注意力机制将多视角空间特征进行聚合(图中蓝色箭头),另一方面通过时序自注意力机制将历史的时序特征融合到当前时刻(图中绿色箭头),最终获得融合时空特征的 BEV 查询向量,组合得到可以支持下游任务的 BEV 视角特征。

BEVFormer 实验介绍

数据集

BEVFormer 使用 nuScenes 数据集进行试验,nuScenes 数据集是第一个提供自动汽车全套传感器数据的大型数据集,包括了 6 个相机、1 个激光雷达、5 个毫米波雷达、以及 GPS 和 IMU。与另一个常用的自动驾驶领域 KITTI 数据集相比,其包含的对象注释多了 7 倍多,nuScenes 自动驾驶数据集因其数据的质量、规模与难度都比之前数据集有大幅提升,而获得了很多研究者的关注。

检测指标

BEVFormer 在 nuScenes 数据集的 上获得了 0.517 的 NDS(nuScenes detection score)、0.416 的 mAP(mean average precision)。

BEVFormer 在 nuScenes 数据集上的指标列表图

训练时间

BEVFormer Base 模型需要每张 GPU 上 28G 的显存,在单张 A100 上训练完成(24 个 epoch)需要大约 14 天的时间,在 8 张 A100 上训练完成需要接近 2 天的时间。

高昂的训练时间成本,阻碍了科研者和开发者在模型改进的快速迭代。因此幻方 AI 将 BEVFormer 适配到幻方萤火集群,帮助使用者快速有效的进行实验迭代。

BEVFormer 萤火训练

训练时间对比

我们首先保持默认参数,只通过增加并行 GPU 数(相当于增大 batch size),进行模型训练时间和模型性能的对比,结果如下面表格所示。从表格中可以看出,随着并行 GPU 数量的增加,训练 24 个 epoch 的时间逐渐缩短,从最长 14 天 20 小时缩短到 3 小时 12 分中。表格中 ~ 符号代表估算时间,即通过单次迭代的时间对总训练时间进行估算。另外,单卡的训练由于时间过长,只进行时间估算,未统计性能。

模型 并行 GPU 数量 训练耗时 NDS mAP
BEVFormer Base 1 ~ 14 天 20 小时 - -
BEVFormer Base 8 ~ 44 小时 24 分 0.517 0.416
BEVFormer Base 80 6 小时 28 分 0.460 0.365
BEVFormer Base 160 3 小时 30 分 0.385 0.337

增加 GPU 数量,速度变快但性能下降?

从上面表格结果可以看出,在相同的参数训练下,随着并行 GPU 数量的增加(相当于 batch size 增大),训练时间越来越短。但是与此同时,随着并行数量的增加,最终训练出的模型性能越来越差,如下图的变化曲线所示。当我们使用 160 张卡进行训练时,训练模型最终的 NDS 和 mAP 指标与 8 卡时大约相差 15 个点。

默认参数性能指标变化曲线图

为什么增加并行 GPU 数(相当于增大 batch size),会导致训练结果变差呢?

我们探究后发现:这是因为深度学习网络的损失函数是非凸函数,基于凸优化理论的优化算法容易收敛到局部最优解。而当我们增大 batch size,梯度下降的方向更加确定,随机性更小,更容易收敛到附近的局部最优解,而无法达到全局最优解。

如何增加并行的同时保持性能?

随着 batch size 的增大,样本梯度的方差会同比减小,梯度方向更确定,如果要保持增大 batch size 之前的梯度方差,我们可以选择增加学习率。

增加学习率有两方面好处,一方面可以提高收敛速率,因为大 batch size 下,梯度估计更准,增大学习率可以加快模型收敛到最优解。另一方面有助于帮助模型逃出局部最优解,增大学习率可以增大下降步幅,使得模型有机会逃出局部最优解,从而收敛到更好的最优解。

那么如何调整到合适的学习率呢?以往针对大 batch size 的研究提出了下面两种缩放学习率的方法:

  • sqrt 缩放:Krizhevsky 等人 [1] 提出,当 batch size 增大 k 倍时,学习率应该增大 sqrt(k) 倍。
  • 线性缩放:Goyal 等人[2] 提出了更简单的线性缩放规则,他们认为当 minibatch 大小乘以 k 时,将学习率乘以 k

我们从实验中发现,按照上述原则调整学习率确实消除了小批量和大批量之间的部分性能差距,但是调整后的模型性能,与小 batch size 的性能还相差甚大。

为什么会有这样的现象呢?文献 [3] 指出,随着 batch size 变得越来越大,得到最佳测试结果所能允许的学习率范围越来越小。也就是说,当 batch size 比较小时,轻微改变学习率对于模型最终的性能影响不大,因此我们可以很容易的找到合适的学习率。但是当 batch size 变大,轻微的学习率变化对最终训练的模型性能影响也随之变大,因此我们需要不断的精细化查找,才能找到合适的参数,这为我们分布式并行训练带来了挑战。

HFai 优化工具实践

为了获得更高的训练精度和更快的训练速度,我们使用幻方 AI 的 Autotune 和 FFRecord 两个工具对训练过程进行优化。

Autotune: 自动搜索集群最佳参数

Autotune 是幻方 AI 针对萤火集群训练平台设计的自动调参工具,可以根据用户设定好的参数范围自动发起集群训练,并统计不同参数下的模型表现。使用时只只需要增加几行代码,就可以使用 Autotune 进行简单的高效的大规模参数调优。

在 BEVFormer 的训练中,Autotune 很好的解决了前面所提到的大规模并行训练时面临的精细化调参的挑战。根据学习率的 sqrt 缩放原则 [1] 和线性缩放原则 [2],我们可以大致确定在使用上百张GPU进行训练时应该使用的学习率的大致范围,通过 Autotune 我们可以发起多个训练任务来同时对指定范围内的参数进行搜索,快速筛选出最好的参数组合。如下图所示,使用 Autotune 我们可以一键发起数个调参任务:

任务截图

这样只需要和正常发起一次训练一样的工作量和等待时间,我们就可以得到调参结果。而在使用 Autotune 完成训练后,我们同样可以简单的查看各参数训练结果的排序,便于筛选最优参数:

结果统计

从上图可见,在使用 160 张 GPU 训练 BEVFormer 时,我们找到的最优参数组合是:lr1e-3warmup_iters600min_lr_ratio1e-4, 在这组参数下我们取得了 0.490 的 NDS 和 0.389 的 mAP 指标,相比默认参数增加并行的结果,NDS 提高了 0.115,相比 8 张 GPU 时也仅仅下降了 0.027,但训练时间缩短了 14 倍。

FFRecord:加速数据读取

前面的文章我们介绍了幻方高性能数据格式 FFRecord,使用 FFRecord 格式我们可以获得更快的数据读取速度。针对 mmlab 系列开源仓库,我们提供了 mmconverter,用户可以通过简单的三步将 mmlab 代码中使用的数据格式转化为 FFRecord 数据格式,达到适配幻方集群文件系统和提高训练速度的目的。

如下表所示,使用 FFRecord 格式后,不同并行数量任务都获得了一定的性能提升,最终 3 小时就可以完成 BEVFormer 的训练。表格中 ~ 符号代表估算时间,即通过单次迭代的时间对总训练时间进行估算。另外,单卡的训练由于时间过长,只进行时间估算,未统计性能。

模型 并行 GPU 数量 训练耗时 训练耗时(FFRecord) NDS mAP
BEVFormer Base 1 ~ 14 天 20 小时 ~ 13 天 12 小时 - -
BEVFormer Base 8 ~ 44 小时 24 分 ~ 39 小时 12 分 0.517 0.416
BEVFormer Base 80 6 小时 28 分 6 小时 0.498 0.391
BEVFormer Base 160 3 小时 12 分 3 小时 10 分 0.490 0.389

下面曲线将使用默认参数的性能曲线与使用 Autotune 搜索到的最佳参数进行了对比。可以看到,使用 Autotune 找到的最佳参数,随着并行数量的增加,性能下降更加平缓,从 1 张 GPU 增加到 160 张 GPU,NDS 和 mAP 都只下降了 0.027,相比默认参数的下降幅度有了很大的提升。

性能变化曲线对比图

看到这里的你有没有遇到过模型训练久、实验效率低的问题呢?快来申请幻方助力科研的免费 GPU 资源,使用幻方一系列深度学习训练神器,提升你的训练体验吧!

参考文献

[1] Krizhevsky, Alex. “One weird trick for parallelizing convolutional neural networks.” arXiv preprint arXiv:1404.5997 (2014).

[2] Goyal, Priya, et al. “Accurate, large minibatch sgd: Training imagenet in 1 hour.” arXiv preprint arXiv:1706.02677 (2017).

[3] Breuel, Thomas M. “The effects of hyperparameters on SGD training of neural networks.” arXiv preprint arXiv:1508.02788 (2015).

[4] 代季峰对话张祥雨 | 自动驾驶感知新时代!新一代环视感知算法BEVFormer有哪些优势。https://hub.baai.ac.cn/view/19213


本文作者: Freja


您可以转载、不违背作品原意地摘录及引用本技术博客的内容,但必须遵守以下条款: 署名 — 您应当署名原作者,但不得以任何方式暗示幻方为您背书,亦不会对幻方的权利造成任何负面影响。 非商业性使用 — 您不得将本技术博客内容用于商业目的。 禁止演绎 — 如果基于该内容改编、转换、或者再创作,您不得公开或分发被修改内容,该内容仅可供个人使用。