混合精度: 4分钟训练ImageNet

在前面三篇里给大家介绍了大批量同步式训练这种近来流行的训练模型的方法，包括LARS优化器和Ring All-reduce架构。今天再来看一些改进，本文的工作是腾讯做出来的，在参考文献[1]。

• 大批量SGD: 1小时训练ImageNet
• LARS: 学习率层次自适应，超大批量32K
• Ring All-reduce: 分布式深度学习的巧妙同步

Overall

论文主要做了三个改进:

• 提出了一种混合精度的训练方法，可以在不损失accuracy的条件下提高单个GPU的通量。
• 提出了一种优化方法可以在超大批量(64k)上进行训练而没有accuracy的损失。
• 提出了一种高度优化的all-reduce算法可以使用1024个P40 GPU在AlexNet和ResNet上分别达到3x和11x的优化比。

挑战

在使用大batch进行模型训练时，通常会遇到两个挑战：

• 尽管大batch可以减少梯度的方差，加快模型训练的速度，但很多研究表明增加到一定程度后会损失最后的结果。
• 在使用大集群时，很难达到接近线性的加速比，因为通信的成本越来越高。

系统设计

先来看一下整个系统的设计，系统主要分为三个部分:

• input pipeline: 在本步训练完成之前，发送下一步训练需要的数据给GPU
• training module: 负责模型的构建和变量管理，在这一部分，模型的前向/后向计算会用mixed precision来做，模型的更新用LARS来做。
• communication module: 使用tensor fusion和混合all-reduce来使得系统在tensor和cluster数目变化的时候具有良好的扩展性。

混合精度

使用半精度FP16，可以在节省显存、网络带宽的同时提高计算的吞吐量。

但现在的大批量计算，一般使用LARS这样的自适应学习率算法，正如我们在上文所讲，LARS是通过计算权重的L2-Norm来得到的，如下面的公式，

这样的话，因为涉及到L2-norm，暴力的将所有32位数值转成16位会导致数字越界，超过16位的上限的问题。为了解决这个问题，使用的混合精度，在前向计算和后向计算的时候，采用FP16位来进行，而在计算LARS的时候，则转为32位进行计算。

混合精度 + LARS是本系统能够具有可扩展性的重要原因。从下图的实验结果中就可以看出。

模型结构提升

论文在AlexNet和ResNet上做了实验，其中在AlexNet上提出了几点改进：

• 去掉bias和Batch Normalization上的权重衰减
• 在合适的地方添加一个Batch Normalization层。

权重衰减是一个常用的用来防止过拟合的手段，其做法就是给损失函数添加一个L2的正则化项，这样在计算梯度的时候就体现为权重衰减。损失函数如下：

梯度如下，最后一项来源于正则化项，就是权重衰减项。

那么，为什么给bias和Batch Normalization去掉正则化项就能带来提升呢？

在神经网络中，BN层的参数和bias占总体参数的比例都非常少，例如，在AlexNet中，BN层参数和bias加起来占总参数量的0.02%。不对它们做正则化并不会带来特别多的方差，反之，做了正则化反而会使得它们欠拟合。

下图显示了去掉BN和bias正则化的效果：

除此之外，还发现AlexNet最后一层pooling过后方差比较大，导致训练比较困难，因此，在最后一个pooling层后添加了一个Batch Normalization层，加完之后上表中的数值可以从57.1%提升到58.8%。

提升通信策略

回顾一下，在Ring All-reduce的分布式环境中，每一步节点的通信量是模型参数总量除以GPU数。因而随着集群的变大，节点间的通信量变少，导致网络带宽没有被完全利用上，从而降低了并行度。

论文提出了两种策略来对此进行提升：

• Tensor Fusion, 在通信的时候，将多个小尺寸的tensor打包成大的，来提高带宽利用率。这里设置一个参数θ，如果tensor小于θ，那么就把tensor放到缓冲池里，当缓冲池大于θ时，进行参数的传递。
• 分层All-reduce, 可以把一些GPU组成一个群，然后先在群上运行All-reduce算法来做梯度同步，然后在不同的群之间进行同步，最后在每个群内部广播梯度。

• 混合All-reduce, 在卷积神经网络中，卷积层参数较少，适合用分层的方式，全连接层参数较多，适合单个GPU。因此，可以把分层All-reduce和普通reduce方式混合起来做梯度的聚合。

实验

实验结果如下图所示，可以看到，加上本文中介绍的这些优化之后，系统可以将训练AlexNet的时间最少降低至4分钟，ResNet-50的训练时间最少降低到6.6分钟。

参考文献

• [1]. Jia, Xianyan, et al. "Highly scalable deep learning training system with mixed-precision: Training imagenet in four minutes." arXiv preprint arXiv:1807.11205 (2018).

由于微信平台算法改版，公号内容将不再以时间排序展示，如果大家想第一时间看到我们的推送，强烈建议星标我们和给我们多点点【在看】。星标具体步骤为：

（1）点击页面最上方"AINLP"，进入公众号主页。

（2）点击右上角的小点点，在弹出页面点击“设为星标”，就可以啦。

感谢支持，比心。

欢迎加入AINLP技术交流群

进群请添加AINLP小助手微信 AINLPer（id: ainlper)，备注NLP技术交流

推荐阅读

这个NLP工具，玩得根本停不下来

征稿启示| 200元稿费+5000DBC（价值20个小时GPU算力）

完结撒花！李宏毅老师深度学习与人类语言处理课程视频及课件（附下载）

从数据到模型，你可能需要1篇详实的pytorch踩坑指南

如何让Bert在finetune小数据集时更“稳”一点

模型压缩实践系列之——bert-of-theseus，一个非常亲民的bert压缩方法

文本自动摘要任务的“不完全”心得总结番外篇——submodular函数优化

Node2Vec 论文+代码笔记

模型压缩实践收尾篇——模型蒸馏以及其他一些技巧实践小结

中文命名实体识别工具（NER）哪家强？

学自然语言处理，其实更应该学好英语

斯坦福大学NLP组Python深度学习自然语言处理工具Stanza试用

关于AINLP

AINLP 是一个有趣有AI的自然语言处理社区，专注于 AI、NLP、机器学习、深度学习、推荐算法等相关技术的分享，主题包括文本摘要、智能问答、聊天机器人、机器翻译、自动生成、知识图谱、预训练模型、推荐系统、计算广告、招聘信息、求职经验分享等，欢迎关注！加技术交流群请添加AINLPer(id：ainlper)，备注工作/研究方向+加群目的。

阅读至此了，分享、点赞、在看三选一吧🙏