论文导读：24分钟完成ImageNet训练，刷新世界纪录

投稿

作者：尤洋
编辑：Natalie

近年来，深度学习的一个瓶颈主要体现在计算上。比如，在一个英伟达的M40 GPU上用ResNet50去训练ImageNet需要14天；如果用一个串行程序在单核CPU上训练可能需要几十年才能完成。这个问题大大限制了科技的发展。为了设计新的模型，研究人员往往需要不断调整模型，再去做实验，每次实验结果都要等一天是无法接受的。所以，在半小时之内完成大型神经网络的训练对科技的进步意义重大，交互式研究能大大提高研发效率。

利用参数服务器实现的异构方法无法保证在大型系统之上稳定起效。而正如Goyal等人于2017年得出的结论，数据并行同步方法对于超大规模深度神经网络（简称DNN）训练而言表现得更为稳定。其基本思路非常简单——在SGD中使用更大的batch size，确保每一次迭代皆可被轻松分布至多处理器处。这里我们考虑以下理想条件。ResNet-50在处理一张225x225像素图片时，需要进行77.2亿次单精度运算。如果我们为ImageNet数据集运行90 epochs的运算，则运算总量为90x128万x77.2亿（10¹⁸）。其实，现在世界上最快的计算机（中国的神威太湖之光）理论上可以每秒完成2*10^17个单精度浮点数计算。所以，如果我们有一个足够好的算法，允许我们充分利用这台超级计算机，应该能够在5秒内完成ResNet-50的训练。

为了实现这一目标，我们需要确保该算法能够使用更多处理器并在每次迭代时加载更多数据——也就是说，要在SGD采用更大的batch size。这里我们用单一英伟达M40 GPU来假定单机器用例。在特定范围内，batch size越大，则单一GPU的速度就越快（如图2所示）。其原因在于，低级矩阵计算库在这种情况下拥有更高的执行效率。在利用ImageNet训练AlexNet模型时，其中每GPU的最优批量大小为512。如果我们希望使用大量GPU并保证每GPU皆拥有理想的执行效率，则应当将批量大小设定为16 x 512 = 8192。

在理想情况下，如果我们固定数据整体访问量并根据处理器数量线性增加batch size，每次迭代的batch size越大，所需要的迭代次数就越少。在增大batch size的同时，我们可以使用更多处理器以保持单循环时间基本维持恒定。如此一来，我们就可以获得线性加速比（如表1所示）。

可惜，SGD算法的batch size并不能无限制地增大。SGD采用较大的batch size时，如果还是使用同样的epochs数量进行运算，则准确度往往低于batch size较小的场景;而且目前还不存在特定算法方案能够帮助我们高效利用较大的batch size。表2所示为基准测试下的目标准确度。举例来说，如果我们将AlexNet的batch size设置为1024以上或者将ResNet-50的batch size设置为8192以上，则准确度的测试结果将严重下降（如表4与图3所示）。

对于大批量训练场景，我们需要确保使用较大的batch size训练能够在相同epochs前提下获得与较小的batch size相近的测试准确度。这里我们之所以保持epochs数量不变，是因为从统计学角度来讲，一个epoch代表着算法与整体数据集接触一次;而从计算角度来讲，固定epochs数量意味着保证浮点运算次数不变。目前业界领先的大批量训练方案包含以下两种技术手段：

(1) 线性缩放（Krizhevsky于2014年提出）：如果我们将batch size由B增加至kB，我们亦需要将学习率由η增加至kη（其中k为倍数）。

(2) 预热模式（Goyal等人于2017年提出）：如果我们使用高学习率（η），则应以较小的η作为起点，而后在前几次epochs中逐步将其递增至较大η。

这里我们使用B、η以及I来分别表示batch size、学习率与迭代次数。如果我们将batch size由B增加到kB，那么迭代次数将由I降低至I/k。这意味着权重的更新频率减少为原本的k分之一。因此，我们可以通过将学习率提升k倍使每一次迭代的更新效率同样提升k倍。预热模式的目标在于避免算法在初始阶段因为使用基于线性缩放的高学习率而无法收敛。利用这些技术手段，研究人员们能够在特定范围之内使用相对较大的batch size（如表3所示）。

然而，我们亦观察到目前业界领先的各类方案的batch szie区间仍然比较有限，其中AlexNet为1024，而ResNet-50为8192。如果我们在ALexNet模型训练当中将batch size增加至4096，则100 epochs情况下的准确度仅能达到53.1%（如表4所示）。我们的目标是在使用更大batch size的情况下至少达到58%的准确度。

据我所知，人们之前还无法有效利用诸如神威太湖之光的超级计算机完成神经网络训练。正如姚期智教授所说的那样，AI的突破口在算法。这个问题的瓶颈就在算法层。我们有足够的计算能力，却无法充分利用。

为了解决大批量训练的准确度问题，我们设计了LARS算法[1]，允许我们将AlexNet的batch size提升到32k（原来只有1024），将ResNet50的batch size也提高到32k。一旦我们完成了batch size的提升，用超级计算机加速神经网络就变得非常简单了。

只需把数据均分到各个机器上，每次循环完成一次对所有机器上参数的均值操作即可。最终，我们用24分钟完成了AlexNet的训练（据我们所知，这是使用AlexaNet进行ImageNet训练的世界最快纪录）。从下面的两个图表可以看出，我们同样可以在一小时内完成ResNet-50的训练。由于LARS使我们得以把batch size扩展到了32k，我们可以使用更为廉价的芯片进行计算。我们选择了Intel KNL（其实就是一种高端CPU），最终总共用了512个KNL，按照市场价的花费是120万美金，远远低于之前Facebook的410万美金 (32台NVIDIA DGX station)。正是大的batch size使我们大大降低了成本。

LARS算法目前已经被用在NVIDIA Caffe和Google Tensorflow之中，NVIDIA Caffe的LARS实现版本已经开源。当然，如果用更大的超级计算机，训练时间还会进一步缩短。

其实一年前，吴恩达教授在国际超算大会上(ISC2016)演讲时就预言，超级计算机将来对AI会非常重要。笔者认为，如果算法的扩展性足够好，超级计算机可以很快完成神经网络训练。那么，神经网络对单机server的需求可能会下降。相反，人们可能会去买超算中心的computational hour。这样，普通人也可以廉价地用上1000台机器。并且超算中心会把这些机器维护好，普通用户不需要自己浪费大量的时间去管理机器和系统软件。当然这对Intel和NVIDIA的芯片厂家是没有影响的，因为他们可以将芯片大批量卖给超算中心。

论文原文链接：
https://arxiv.org/pdf/1709.05011.pdf

作者介绍：
尤洋，UC Berkeley 计算机科学博士在读，LARS算法开发者之一。
http://www.cs.berkeley.edu/~youyang/