ImageNet Training in Minutes

摘要与介绍

目前，缩短 DNN 的训练时间是一个热点，通用的方法聚焦于开发一种新的训练算法可在保证不损失精确性的前提下增大 batch 的大小，当然这些方法基本都基于数据并行的同步 SGD 进行参数更新。这其中经历了 batch_size 从 1K 到 8K，再增大到 32K 的过程。本文通过 LARS 算法使用 1000 个 cpu 在 11 分钟内完成了 100 个 epoch 的 Imagenet 训练集的训练，模型是 AlexNet，batch 设的是 32K，取的了 58.6%的 accuracy。另外在 resnet50 模型上，通过 64 个 epoch 的训练，在 14 分钟之内取得了 74.9% top-1 test 集的 accuracy 。另外，当 batch_size 大于 16K 时，通过 LARS 算法训练的准确率要高于 Facebook 的一小时训练 imagenet 那篇文章中的结果。如下图：

本文提出的算法使用了两个硬件：Intel Skylake CPU 和 Intel KNL（通用的深度学习加速器）。看到这心凉了，又是硬件在支撑。。。

计算通信比是指模型的计算代价/通信代价，一般计算通信比越高，越容易通过分布式来进行加速。如下表：ResNet50 的计算通信比要高于 AlexNet，所以 ResNet50 的分布式训练加速效果要好。另外大的 batch_size 可以减小通信代价，因为大的 batch 意味着更少轮次的迭代。

本文使用的方法是：LARS + Warmup。

下面具体介绍下 LARS 算法：

标准的 SGD 对于网络的每一层都使用相同的学习率 LR，当学习率很大时，更新的幅度会很大，容易导致发散的情况。这就使得模型在最初阶段对于 weight 的初始化和学习率设置显得十分重要。另外权重与梯度的 L2-norm 比 在不同 weight、layer 见变动很大。

warm-up 方法会在一开始先从小的学习率开始，然后逐步增大到大的学习率。LARS 则会对于每一层使用一个 local 学习率 [公式] ：

其中 [公式] 是整体的学习率，[公式]是每一层的学习率，计算方法为：

如果加上 weight decay 参数 [公式] 后，上式可写成：

完整的 LARS 算法为：

实验结果：

这里就贴一个 resnet50 的结果吧，可以看出使用 LARS 算法可以将 batch size 扩展到 32k，另外还可以使用更廉价的芯片进行计算。使用 512 个 Intel KNL 可在 1 个小时内完成训练。512 个 KNL 按照市场价是 120 万美金，远远低于之前 Facebook 的 410 万美金 (32 台 NVIDIA DGX 工作站)。

与其它方法的对比：

参考文献：

https://arxiv.org/pdf/1709.05011.pdf

https://arxiv.org/pdf/1708.03888.pdf

本文转载自 Alex-zhai 知乎账号。

原文链接：https://zhuanlan.zhihu.com/p/81243154