[论文解读] MegDet: A Large Mini-Batch Object Detector

Chao Peng, Tete Xiao+6 more|arXiv (Cornell University)|Nov 20, 2017

Advanced Neural Network Applications参考文献 34被引用 29

一句话总结

MegDet 提出了一种大规模小批量目标检测框架，可在使用 128 张 GPU 的情况下将小批量大小提升至 256，实现 COCO 2017 上 52.5 mmAP 的性能，并将训练时间从 33 小时减少至仅 4.1 小时。该方法引入了预热学习率策略和跨 GPU 批归一化（CGBN），以在大规模训练中稳定训练过程并提升精度。

ABSTRACT

The improvements in recent CNN-based object detection works, from R-CNN [11], Fast/Faster R-CNN [10, 31] to recent Mask R-CNN [14] and RetinaNet [24], mainly come from new network, new framework, or novel loss design. But mini-batch size, a key factor in the training, has not been well studied. In this paper, we propose a Large MiniBatch Object Detector (MegDet) to enable the training with much larger mini-batch size than before (e.g. from 16 to 256), so that we can effectively utilize multiple GPUs (up to 128 in our experiments) to significantly shorten the training time. Technically, we suggest a learning rate policy and Cross-GPU Batch Normalization, which together allow us to successfully train a large mini-batch detector in much less time (e.g., from 33 hours to 4 hours), and achieve even better accuracy. The MegDet is the backbone of our submission (mmAP 52.5%) to COCO 2017 Challenge, where we won the 1st place of Detection task.

研究动机与目标

• 解决小批量训练在目标检测中效率低下且不准确的问题，该问题限制了训练速度和批归一化统计量的稳定性。
• 克服大规模小批量检测器训练的挑战，其中过大的学习率会导致不稳定，而过小的学习率则会损害收敛性。
• 通过将小批量大小扩展至 256 和 GPU 数量扩展至 128，实现在不损失精度的前提下加快训练周期。
• 通过引入跨 GPU 批归一化（CGBN）来改进目标检测中的批归一化统计量。

提出的方法

• 引入一种预热学习率策略，逐步提高初始学习率，以在大规模小批量训练中稳定训练过程。
• 提出跨 GPU 批归一化（CGBN），通过聚合多张 GPU 上的批量统计量来提升 BN 的精度和训练稳定性。
• 应用线性缩放规则，并基于目标检测中保持损失方差一致性的新解释进行调整。
• 采用长周期训练策略并结合学习率衰减，以提升大规模小批量设置下的收敛性和最终精度。
• 当 BN 统计量大小超过 32 时，实现 BN 统计量的次线性内存使用，以减少计算开销。
• 将 CGBN 与 OHEM、空洞卷积、多尺度训练以及更强的主干网络结合，以提升模型性能。

实验结果

研究问题

• RQ1 大规模小批量训练能否在目标检测中成功应用，实现精度损失最小化和训练速度显著提升？
• RQ2 与标准 BN 相比，跨 GPU 批归一化（CGBN）在大规模小批量设置下如何提升训练稳定性和精度？
• RQ3 在 COCO 上进行大规模小批量目标检测时，最优的小批量大小和 BN 统计量大小分别是什么？
• RQ4 预热学习率策略是否能有效稳定目标检测中大规模小批量的训练过程？
• RQ5 与标准的小批量训练相比，大规模小批量训练在收敛速度和最终精度方面表现如何？

主要发现

• 在 128 张 GPU 上使用 256 的小批量大小训练 MegDet，在 COCO 2017 测试开发集上达到 52.5 mmAP，获得目标检测挑战赛第一名。
• 训练时间从基线（在 8 张 GPU 上使用 16 的小批量大小）的 33.2 小时缩短至 4.1 小时（在 128 张 GPU 上使用 256 的小批量大小），实现 8 倍加速。
• 最佳性能出现在 BN 统计量大小为 32 时；更小的尺寸（如 2、4、8）因统计量质量差而降低精度，而更大的尺寸（如 64）也会导致性能下降。
• 预热学习率策略使大规模小批量训练得以稳定进行，即使在高学习率下也能实现收敛。
• CGBN 显著提升了训练稳定性和精度，尤其在大规模小批量设置下，且能在不同小批量大小下保持一致的性能表现。
• 更长的训练策略（如 256（长））可略微提升精度（37.7 vs. 37.1 mmAP），并在大规模小批量设置下表现出一致的收敛性。