2022年的当下，混合精度 (Automatically Mixed Precision, AMP) 训练已经成为了炼丹师的标配工具，仅仅只需几行代码，就能让显存占用减半，训练速度加倍。
AMP 技术是由百度和 NIVDIA 团队在 2017 年提出的 (Mixed Precision Training)，该成果发表在 ICLR 上。PyTorch 1.6之前，大家都是用 NVIDIA 的 apex 库来实现 AMP 训练。1.6 版本之后，PyTorch 出厂自带 AMP。
这篇文章由浅入深地讲解了: 如何在 PyTorch 中使用 AMP、AMP 的原理、AMP 的代码实现。
1. 如何在 PyTorch 中使用 AMP
如果你是新手，只是想简单地试用一下 AMP，只需要将相关训练代码
output = net(input)
loss = loss_fn(output, target)
loss.backward()
optimizer.step()
optimizer.zero_grad()修改如下即可。
with torch.cuda.amp.autocast():
output = net(input)
loss = loss_fn(output, target)
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()
optimizer.zero_grad()如果 GPU 支持 Tensor Core (Volta、Turing、Ampere架构)，AMP 将大幅减少显存消耗，加快训练速度。对于其它类型的 GPU，仍可以降低显存，但训练速度可能会变慢。
2. AMP的原理
2.1 什么是FP16
半精度浮点数 (FP16)是一种计算机使用的二进制浮点数数据类型，使用 2 字节 (16 位) 存储，表示范围为 ![]() 。而 PyTorch 默认使用单精度浮点数 (FP32)来进行网络模型的计算和权重存储。FP32 在内存中用 4 字节 (32 位) 存储，表示范围为 ![]() 。可以看到 FP32 能够表示的范围要比 FP16 大的多得多。
此外浮点数还有一个神奇的特点: 当两个数字相差太大时，相加是无效的，又称舍入误差。
用一段代码来展示
>>> # FP32相加不会有问题。
>>> torch.tensor(2**-3) + torch.tensor(2**-14)
tensor(0.1251)
>>> # FP16相加，较小的数字会被忽略掉。因为在[2**-3, 2**-2]内，FP16表示的固定间隔为2**-13。
>>> # 也就是说比2**-3大的下一个数字为2**-3 + 2**-13，因此2**-14加了跟没加一样。
>>> # half()的作用是将FP32转化为FP16。
>>> torch.tensor(2**-3).half() + torch.tensor(2**-14).half()
tensor(0.1250, dtype=torch.float16)
>>> # 将2**-14换成2**-13就可以了。
>>> torch.tensor(2**-3).half() + torch.tensor(2**-13).half()
tensor(0.1251, dtype=torch.float16)2.2 为什么要用 FP16
如果我们在训练过程中将 FP32 替代为 FP16，有以下两个好处:
- 减少显存占用: FP16 的显存占用只有 FP32 的一半，这使得我们可以用更大的 batch size。
- 加速训练: 使用 FP16，模型的训练速度几乎可以提升1倍。
2.3 为什么只用 FP16 会有问题
如果我们简单地把模型权重和输入从 FP32 转化成 FP16，虽然速度可以翻倍，但是模型的精度会被严重影响。原因如下:
- 上/下溢出: FP16 的表示范围不大，超过
![]() 的数字会上溢出变成 inf，小于 ![]() 的数字会下溢出变成 0。下溢出更加常见，因为在网络训练的后期，模型的梯度往往很小，甚至会小于 FP16 的下限![]() ，此时梯度值就会变成 0，模型参数无法更新。下图为 SSD 网络在训练过程中的梯度统计，有 67% 的值下溢出变成 0。
![]()
- 舍入误差: 就算梯度不会上/下溢出，如果梯度值和模型的参数值相差太远，也会发生舍入误差的问题。假设模型参数
![]() ，学习率 ![]() ，梯度 ![]() ， ![]()
2.4 解决方案
为了解决下溢出的问题，论文中对计算出来的 loss 值进行缩放 (scale)，由于链式法则的存在，对 loss 的缩放会作用在每个梯度上。缩放后的梯度，就会平移到 FP16 的有效范围内。这样就可以用 FP16 存储梯度而又不会溢出了。此外，在进行更新之前，需要先将缩放后的梯度转化为 FP32，再将梯度反缩放 (unscale) 回去。
注意这里一定要先转成 FP32，不然 unscale 的时候还是会下溢出。 缩放因子 (loss_scale) 一般都是框架自动确定的，只要没有发生 inf 或者 nan，loss_scale 越大越好。因为随着训练的进行，网络的梯度会越来越小，更大的 loss_scale 可以更加充分地利用 FP16 的表示范围。
为了实现 FP16 的训练，我们需要把模型权重和输入数据都转成 FP16，反向传播的时候就会得到 FP16 的梯度。如果此时直接进行更新，因为梯度 * 学习率的值往往较小，和模型权重的差距会很大，可能会出现舍入误差的问题。
所以解决思路是: 将模型权重、激活值、梯度等数据用 FP16 来存储，同时维护一份 FP32 的模型权重副本用于更新。在反向传播得到 FP16 的梯度以后，将其转化成 FP32 并 unscale，最后更新 FP32 的模型权重。因为整个更新过程是在 FP32 的环境中进行的，所以不会出现舍入误差。
FP32 权重备份解决了反向传播的舍入误差问题。
![]()
对于那些在 FP16 环境中运行不稳定的模块，我们会将其添加到黑名单中，强制它在 FP32 的精度下运行。比如需要计算 batch 均值的 BN 层就应该在 FP32 下运行，否则会发生舍入误差。还有一些函数对于算法精度要求很高，比如 torch.acos()，也应该在 FP32 下运行。论文中的黑名单只包含 BN 层。
如何保证黑名单模块在 FP32 环境中运行: 以 BN 层为例，将其权重转为 FP32，并且将输入从 FP16 转成 FP32，这样就可以保证整个模块是在 FP32 下运行的。 黑名单解决了某些函数在 FP16 环境下的算术不稳定的问题。
![]()
Tensor Core 可以让 FP16 做矩阵相乘，然后把结果累加到 FP32 的矩阵中。这样既可以享受 FP16 高速的矩阵乘法，又可以利用 FP32 来消除舍入误差。
搞不懂 Tensor Core 是如何应用到 AMP 中的。有人说 Tensor Core 可以帮助我们利用 FP16 的梯度来更新 FP32 的模型权重。但是阅读了 apex 的源码之后，我发现 FP16 的梯度会先转化为 FP32，再做更新，所以权重更新和 Tensor Core 并无关系。以后弄明白了再回来补充吧。 2.5 一些思考
其实将 FP16 和 FP32 混合起来使用是必然的结果，有以下几个原因
- 在网络训练的后期，梯度的值非常小，可能会让 FP16 下溢出。如果不使用 FP32，即使我们通过缩放操作暂时规避了这个问题，权重更新时的 unscale 操作还是会让梯度下溢出。
- 承接第 1 条，就算梯度能够以 FP16 表示，但是梯度 * 学习率可能会下溢出。所以权重更新这步操作还是得在 FP32 下运行。
- 承接第 2 条，就算梯度 * 学习率不会下溢出，其值相对于权重本身也是非常小的。权重 + 梯度 * 学习率这步操作中可能会发生舍入误差的问题。
- 承接第 3 条，就算权重 + 梯度 * 学习率不会发生舍入误差，有些算子在 FP16 下也是不稳定的，比如 BN、torch.acos 等。
3. NVIDIA apex库代码解读
首先介绍下 apex 提供的几种 opt-level: o1, o2, o3, o4。注意这里是字母"o"不是数字"0"。
![]()
图片来自: 全网最全-混合精度训练原理 o0是纯FP32，用来当精度的基准。o3是纯 FP16，用来当速度的基准。
重点讲 o1 和 o2 。我们之前讲的 AMP 策略其实就是 o2: 除了 BN 层的权重和输入使用 FP32，模型的其余权重和输入都会转化为 FP16。此外还会创建一个 FP32 的权重副本来执行更新操作。
和 o2 不同， o1 不再需要 FP32 权重备份，因为 o1 的模型一直都是 FP32。 可能有些读者会好奇，既然模型参数是 FP32，那怎么在训练过程中使用 FP16 呢？答案是 o1 建立了一个 PyTorch 函数的黑白名单，对于白名单上的函数，强制要求其用 FP16，即会将函数的参数先转化为 FP16，再执行函数本身。黑名单则强制要求 FP32。以 nn.Linear 为例， 这个模块有两个权重参数 weight 和 bias，输入为 input，前向传播就是调用了 torch.nn.functional.linear(input, weight, bias)。 o1 模式会将 input、weight、bias先转化为 FP16 格式 input_fp16、weight_fp16、bias_fp16，再调用函数 torch.nn.functional.linear(input_fp16, weight_fp16, bias_fp16)。这样一来就实现了模型参数是 FP32，但是仍然可以使用 FP16 来加速训练。
o1 还有一个细节: 虽然白名单上的 PyTorch 函数是以 FP16 运行的，但是产生的梯度是 FP32，所以不需要手动将其转成 FP32 再 unscale，直接 unscale 即可。
个人猜测 PyTorch 会让每个 Tensor 本身的数据类型和梯度的数据类型保持一致，虽然产生了 FP16 的梯度，但是因为权重本身是 FP32，所以框架会将梯度也转化为 FP32。 如果说o1是 FP16 + FP32，更加激进的o2就是 almost FP16 (几乎全是 FP16)。通常来说 o1 比 o2 更稳，一般先选择 o1，再尝试 o2 看是否掉点，如果不掉点就用 o2。
3.1 apex 的 o1 实现
- 根据黑白名单对 PyTorch 内置的函数进行包装。白名单函数强制 FP16，黑名单函数强制 FP32。其余函数则根据参数类型自动判断，如果参数都是 FP16，则以 FP16 运行，如果有一个参数为 FP32，则以 FP32 运行。
- 将 loss_scale 初始化为一个很大的值。
- 对于每次迭代
(a). 前向传播: 模型权重是 FP32，按照黑白名单自动选择算子精度。
(b). 将 loss 乘以 loss_scale
(c). 反向传播，因为模型权重是 FP32，所以即使函数以 FP16 运行，也会得到 FP32 的梯度。
(d). 将梯度 unscale，即除以 loss_scale
(e). 如果检测到 inf 或 nan
i. loss_scale /= 2
ii. 跳过此次更新
(f). optimizer.step()，执行此次更新
(g). 如果连续2000次迭代都没有出现 inf 或 nan，则 loss_scale *= 2
3.2 apex 的 o2 实现
- 将除了 BN 层以外的模型权重转化为 FP16，并且包装了 forward 函数，将其参数也转化为 FP16
- 维护一个 FP32 的模型权重副本用于更新
- 将 loss_scale 初始化为一个很大的值。
- 对于每次迭代
(a). 前向传播: 除了 BN 层是 FP32，模型其它部分都是 FP16。
(b). 将 loss 乘以 loss_scale
(c). 反向传播，得到 FP16 的梯度
(d). 将 FP16 梯度转化为 FP32，并 unscale
(e). 如果检测到 inf 或 nan
i. loss_scale /= 2
ii. 跳过此次更新
(f). optimizer.step()，执行此次更新
(g). 如果连续2000次迭代都没有出现 inf 或 nan，则 loss_scale *= 2
此外，还推荐阅读 MMCV 对于 AMP 的 o2 实现，代码比 apex 更加清晰。但因为我想同时讲 o1 和 o2，就没有选择解读 MMCV 的代码，有兴趣的读者可以进一步研究。
4. 参考
Train With Mixed Precision
AUTOMATIC MIXED PRECISION PACKAGE - TORCH.CUDA.AMP
【PyTorch】唯快不破：基于Apex的混合精度加速
浅谈混合精度训练
全网最全-混合精度训练原理 |
|
转播
