12分钟训练COCO模型!速度精度双提高

【摘要】中科院自动化所图像与视频分析组提出一种基于周期性矩衰减优化的大规模目标检测快速训练方法。该方法在保持精度的前提下,充分利用大规模集群的计算优势,成功将目标检测任务的batch size扩大到1000以上,大幅度地减少训练时间。同时,该方法在大规模检测数据集OpenImages上也带来了10个点以上的性能提升。

该文已被ECCV 2020 收录。

近日,中科院自动化研究所图像与视频分析组提出了一种针对目标检测任务的大批次(large batch)快速训练的优化算法Periodical Moments Decay LAMB(PMD-LAMB,周期性矩衰减的LAMB算法)。该算法首次提出对LAMB算法[1]中的累积一阶矩和二阶矩实现周期性衰减,从根本上摆脱了滞后的历史梯度对当前网络优化的负面影响,有效加快神经网络的收敛速度

在此基础上,本文提出了一个适用于目标检测任务的大批次目标检测训练框架LargeDet,突破了当前检测器只能用小batch size的训练约束,第一次将目标检测任务的训练batch size扩大到1056并且保持网络的收敛性。

通过使用大batch size训练,该方法首次将COCO的检测训练时长缩短至12min,将实例分割(Instance Segmentation)的训练时长缩短至17min,真正实现检测模型的分钟级迭代。同时,该方法在大规模检测数据集OpenImages上也有大幅度的性能涨点。

研究背景

近年来,深度学习遇到的一大挑战就在于模型的训练时间上。比如,在一个英伟达的 M40 GPU 上用 ResNet50 去训练 ImageNet 需要 14 天;如果用一个串行程序在单核 CPU 上训练可能需要几十年才能完成,这个问题的凸显在某种程度上限制着深度学习的发展。在新的模型设计过程中,研究人员需要不断地调整模型,将调整过的模型拿去做实验,然而每次的实验结果都需要等待几天到几十天左右,这在一定程度上拉低了设计效率。

所以,在几个小时之内完成大型神经网络的训练对节省时间和提升效率具有重要意义。

目标检测是智能视觉感知与理解的基础,目前被广泛应用在现实生活中的各个领域,如智能视频监控、自动驾驶、工业检测,无人超市等。目标检测在于检测图片中所有感兴趣类别的物体,并给出尽可能准确的包围框,从而加强了对模型设计的要求,同时也加重了模型训练的难度。

而目前,数据体量不断增大,比如,OpenImages图像集包含了超过3000万张图片和1500万个边界框(标签),那是18TB的图像数据!传统的小batch size训练方法往往要耗费至少几周的时间才能完成对一个检测模型的训练,如使用单机高配8块V100 GPU训练Res50-FPN和 Res101-FPN,分别需要3.6天和4.5天,这极大地延长了算法的迭代周期,不利于快速探索各种新技术。

因此,设计针对目标检测任务的大batch size快速训练方法至关重要,有助于推动AI算法的持续创新与落地

如何保证网络在大规模训练下不发散并且同时保持精度?

在网络训练过程中,学习率设置的经验法则通常是学习率线性尺度原则(LSR),即:当训练batch size扩大K倍时,学习率也要相应的扩大K倍以获得相当的精度。但是,当学习率扩大到一定程度时,网络便会发散。如果使用一个比较小的学习率,网络可能不会发散,但是会面临精度下降的问题。因此,如何保证网络在大batch size训练下不发散并且同时保持精度是一个研究的难点。

MegDet[2]是第一个研究目标检测任务的大batch size训练的工作。通过利用学习率线性尺度原则、学习率预热(Warm up)以及Cross-GPU BatchNorm,MegDet成功将目标检测任务中的训练batch size提升到256,但是当batch size大于256时能否保证检测器的收敛,这个问题仍然未知。因此,本工作对于更大batch size下检测器的训练进行了探究,并给出了可行的解决方案

方法简述

按照学习率线性尺度原则,当训练batch size设置的比较大时,学习率也要设置的比较大以获得与小batch size训练相同的精度,但是过大的学习率会导致网络的发散。这是因为网络中的各层对于学习率的敏感程度不同。有的层对于学习率比较敏感,较大的学习率就会导致该层的梯度过量更新,从而导致整个网络的发散。

如图1所示为使用学习率线性尺度原则时用不同batch size训练Res50-FPN检测器的精度。当batch size小于等于256时,学习率线性尺度原则可以保持与baseline(batch size为16)几乎相同的精度水平。

然而,当batch size大于256时,按照比例增加学习率将导致网络性能下降甚至发散,此时LSR已不再有效。如若使用一个比较小的学习率(正常学习率的一半),网络虽然可以收敛,但是会有很大的精度损失。

图1 Res50-FPN使用学习率线性尺度原则在不同batch size训练时的精度

(LSR表示学习率线性尺度原则,half表示学习率变为原来的一半)

为了解决大batch size训练时网络发散的问题,LAMB算法使用神经网络每一层的参数的模值和梯度更新步的模值的比值来对每一层的学习率进行缩放,从而使得网络的各层有比较合适的学习率。

LAMB算法虽然在一定程度上保证了网络在较大的学习率下的收敛性,但是其收敛速度比较慢,不能满足在有限的迭代次数达到比较高的精度的要求。在网络的训练过程中,神经网络变得越来越智能,其精度越来越高,产生的梯度也越来越能反映正确的梯度优化方向。

但是LAMB算法每一次对网络进行参数更新时,其更新量比较依赖于累积的历史梯度,滞后的历史梯度会阻碍网络的快速收敛。

如公式(1)(2)所示,为t时刻网络产生的梯度,分别为梯度的一阶矩和二阶矩,分别为一阶矩和二阶矩的滑动平均系数。公式(3)为网络的当前更新步的计算公式。由于通常是比较大的值(默认分别为0.9,0.999),受历史梯度的影响比较大,从而导致当前更新步也主要依赖于历史梯度。但是随着网络的训练,历史梯度并不能准确的反映最优的优化方向。因此,减小历史梯度对当前更新步的负面影响可以加快网络的收敛。

PMD-LAMB算法引入一个周期性衰减的系数来对累积的梯度一阶矩和二阶矩进行衰减,如公式(4)(5)(6)所示。在每个周期T内,系数从1衰减到0(如图2所示),累积的一阶和二阶动量对当前的更新步的贡献越来越小,从而缓解历史梯度对于当前更新步的负面影响,使网络能够更快的收敛到最优点。如图3所示为LAMB算法和PMD-LAMB算法的loss曲线对比,从图中可以看出PMD-LAMB相对于LAMB算法收敛得更快,有更小的loss值。PMD-LAMB算法的流程图如图4所示。

图2 变化曲线(T=200)

图3 LAMB和PMD-LAMB的loss曲线对比

以PMD-LAMB算法为核心,本文提出了适用于目标检测任务的大batch size训练算法LargeDet。LargeDet由三部分组成:学习率线性尺度原则,同步batch归一化(Sync BN)以及PMD-LAMB算法。

与此同时,本文提出了三条指导原则用来指导网络在不同batch size下的训练:

1. 学习率线性尺度原则在batch size小于等于256时仍适用,但是当batch size大于256时会导致网络精度严重下降或发散;

2. Sync BN在大batch size训练时可以获得更准确的BN统计量,从而加速网络的收敛;

3. PMD-LAMB保证了检测器在更大batch size训练时的收敛性,比LAMB具有更快的收敛速度。

图4. PMD-LAMB算法

实验验证

我们在常用目标检测数据集COCO和大规模检测数据集OpenImages上进行了实验。如表1所示,该方法首次将检测的batch size扩展到1056。由于batch size为1056时,总的迭代次数过少,无法获得准确的BN统计值,所以适当的延长迭代次数可以提升性能。

表1. COCO上的对比实验

(1.4x表示1.4倍的迭代次数)

如表2所示,通过160块NVIDIA Tesla V100 GPU的并行训练,本方法:

将ResNet50-FPN在COCO上的训练时间缩短到12分钟以内(精度标准为[email protected]:0.95=36.7);

将ResNet50-MaskRCNN在COCO上的训练时间缩短到17分钟以内(精度标准为box [email protected]:0.95=37.6, segment [email protected]:0.95=33.9)。

表2. 快速训练COCO的实验

(Mask表示Mask RCNN)

本方法在大规模检测数据集Open Images上也进行了实验,如表3所示。本方法不仅可以充分利用计算资源,线性地减少训练时间,还可以获得显著的精度提升。

表3. Open Images上的实验结果(600类)

论文题名:

Large Batch Optimization for Object Detection: Training COCO in 12 Minutes" (CV君:目前论文还没有公布)

参考文献

[1] Tong Wang, Yousong Zhu, Chaoyang Zhao, Wei Zeng, Yaowei Wang, Jinqiao Wang, and Ming Tang. Large Batch Optimization for Object Detection: Training COCO in 12 Minutes. ECCV 2020.

[2] You, Y., Li, J., Hseu, J., Song, X., Demmel, J., Hsieh, C.: Reducing BERT pretraining time from 3 days to 76 minutes. CoRR abs/1904.00962 (2019).

[3] Peng, C., Xiao, T., Li, Z., Jiang, Y., Zhang, X., Jia, K., Yu, G., Sun, J.: Megdet: A large mini-batch object detector. In: 2018 IEEE Conference on Computer Vision and Pattern Recognition, CVPR 2018, Salt Lake City, UT, USA, June 18-22, 2018. pp. 6181–6189 (2018).

备注:目标检测

目标检测交流群

2D、3D目标检测等最新资讯,若已为CV君其他账号好友请直接私信。

我爱计算机视觉

微信号:aicvml

QQ群:805388940

微博知乎:@我爱计算机视觉

投稿:[email protected]

网站:www.52cv.net

在看,让更多人看到