训练提速60%！只需5行代码，PyTorch 1.6即将原生支持自动混合精度训练

PyTorch 1.6 nightly增加了一个子模块 amp，支持自动混合精度训练。值得期待。来看看性能如何，相比Nvidia Apex 有哪些优势？

即将在 PyTorch 1.6上发布的 torch.cuda.amp 混合精度训练模块实现了它的承诺，只需增加几行新代码就可以提高大型模型训练50-60% 的速度。

预计将在 PyTorch 1.6中推出的最令人兴奋的附加功能之一是对自动混合精度训练（automatic mixed-precision training）的支持。

混合精度训练是一种通过在半精度浮点数 fp16上执行尽可能多的操作来大幅度减少神经网络训练时间的技术，fp16 取代了PyTorch默认的单精度浮点数 fp32。最新一代 NVIDIA GPU 搭载了专门为快速 fp16矩阵运算设计的特殊用途张量核（tensor cores）。

然而，到目前为止，这些张量核仍然很难用，因为它需要手动将精度降低的操作写入模型中。这就是自动化混合精度训练的用武之地。即将发布的 torc h.cuda.amp API 将允许你只用五行代码就可以在训练脚本中实现混合精度训练！

混合精度是如何工作的

在我们理解混合精度训练是如何工作的之前，首先需要回顾一下浮点数。

在计算机工程中，像1.0151或566132.8这样的十进制数传统上被表示为浮点数。由于我们可以有无限精确的数字(想象一下π) ，但存储它们的空间是有限的，我们必须在精确度(在舍入数字前，我们可以在数字中包含的小数的数量)和大小(我们用来存储数字的位数)之间做出妥协。

浮点数的技术标准 IEEE 754设定了以下标准：fp64, 又名双精度或"double" ，最大舍入误差 ~ 2^-52fp32, 又名单精度或"single"，最大舍入误差 ~ 2 ^-23fp16, 又名半精度或"half" ，最大舍入误差 ~ 2 ^-10。

Python float 类型为 fp64，而对内存更敏感的PyTorch 使用 fp32作为默认的 dtype。

混合精度训练的基本思想很简单: 精度减半(fp32→ fp16) ，训练时间减半。

最困难的是如何安全地做到这一点。

注意，浮点数越小，引起的舍入误差就越大。对“足够小“的浮点数执行的任何操作都会将该值四舍五入到零！这就是所谓的underflowing，这是一个问题，因为在反向传播中很多甚至大多数梯度更新值都非常小，但不为零。在反向传播中舍入误差累积可以把这些数字变成0或者 nans; 这会导致不准确的梯度更新，影响你的网络收敛。

2018年ICLR论文 Mixed Precision Training 发现，简单的在每个地方使用 fp16 会“吞掉”梯度更新小于2^-24的值——大约占他们的示例网络所有梯度更新的5% :

混合精度训练是一套技术，它允许你使用 fp16，而不会导致你的模型训练发生发散。这是三种不同技术的结合。

第一，维护两个权重矩阵的副本，一个“主副本”用 fp32，一个半精度副本用 fp16。梯度更新使用 fp16矩阵计算，但更新于 fp32矩阵。这使得应用梯度更新更加安全。

第二，不同的向量操作以不同的速度累积误差，因此要区别对待它们。有些操作在 fp16中总是安全的，而其它操作只在 fp32中是可靠的。与其用 fp16跑整个神经网络，不如一些用半精度另外的用单精度。这种 dtypes 的混合就是为什么这种技术被称为“混合精度”。

第三，使用损失缩放。损失缩放是指在执行反向传播之前，将损失函数的输出乘以某个标量数(论文建议从8开始)。乘性增加的损失值产生乘性增加的梯度更新值，“提升”许多梯度更新值到超过fp16的安全阈值2^-24。只要确保在应用梯度更新之前撤消缩放，并且不要选择一个太大的缩放以至于产生 inf 权重更新(overflowing) ，从而导致网络向相反的方向发散。

将这三种技术结合在一起，作者可以在显著加速的时间内训练好多种网络以达到收敛。至于benchmarks，我建议读一读这篇只有9页的论文！

张量核（tensor cores）是如何工作的

虽然混合精度训练节省内存(fp16矩阵只有 fp32矩阵的一半大小) ，但如果没有特殊的 GPU 支持，它并不能加速模型训练。芯片上需要有可以加速半精度操作的东西。在最近几代 NVIDIA GPU中这东西叫: 张量核。

张量核是一种新型的处理单元，针对一个非常特殊的操作进行了优化: 将两个4 × 4 fp16矩阵相乘，然后将结果加到第三个4 × 4 fp16或 fp32矩阵(一个“融合乘法加（fused multiply add）”)中。

更大的 fp16 矩阵乘法操作可以使用这个操作作为他们的基本构件来实现。由于大多数反向传播都可以归结为矩阵乘法，张量核适用于网络中几乎任何计算密集层。

陷阱: 输入矩阵必须是 fp16。 如果你正在使用带有张量核的 GPU 进行训练，而没有使用混合精度训练，你不可能从你的显卡中得到100% 的回报！ 在 fp32中定义的标准 PyTorch 模型永远不会将任何 fp16数学运算应用到芯片上，因此所有这些极其强悍的张量核都将处于空闲状态。

张量核在2017年末在上一代Volta体系结构中被引入，当代Turing有了一些改进，并将在即将推出的Ampere中看到进一步的改进。云上通常可用的两款GPU 是 V100(5120个 CUDA 核，600个张量核)和 T4(2560个 CUDA 核，320个张量核)。

另一个值得记住的难题是firmware。尽管 CUDA 7.0或更高版本都支持张量核操作，但早期的实现据说有很多 bug，所以使用 CUDA 10.0或更高版本很重要。

Pytorch 自动混合精度是如何工作的

有了这些重要的背景知识，我们终于可以开始深入研究新的 PyTorch amp API 了。

混合精度训练在技术上已经永远成为可能: 手动运行部分网络在 fp16中，并自己实现损失缩放。自动混合精度训练中令人兴奋的是“自动”部分。只需要学习几个新的 API 基本类型: torch.cuda.amp.GradScalar 和 torch.cuda.amp.autocast。启用混合精度训练就像在你的训练脚本中插入正确的位置一样简单！

为了演示，下面是使用混合精度训练的网络训练循环的一段代码。# NEW标记定位了增加了新代码的地方。

（编辑：辽源站长网）

【声明】本站内容均来自网络，其相关言论仅代表作者个人观点，不代表本站立场。若无意侵犯到您的权利，请及时与联系站长删除相关内容!