serendipity
同济大学
行人搜索
2022 混合精度合精度 (Automatically Mixed Precision, AMP) 训练已经成为炼丹师的标准工具，显存占用可以减半，训练速度可以加倍。
AMP 百度和技术 NIVDIA 团队在 2017 年提出的 (Mixed Precision Training [1])这一结果发表在 ICLR 上。PyTorch 1.6之前，大家都是用 NVIDIA 的 apex [2] 库来实现 AMP 训练。1.6 版本之后，PyTorch 出厂自带 AMP。
本文由浅入深解释：
如果你是新手，你只想简单尝试一下 AMP，相关的训练代码只需要
output=net(input) loss=loss_fn(output,target) loss.backward() optimizer.step() optimizer.zero_grad()修改如下。
withtorch.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 字节 (16 位) 存储，表示范围为 。而 PyTorch 默认使用计算网络模型和存储权重。FP32 在内存中用 4 字节 (32 位) 存储，表示范围为
402 Payment Required
。可以看到 FP32 能表示的范围比 FP16 大的多得多。此外，浮点数还有一个神奇的特征：当两个数字之间的差异太大时，加起来是无效的，也称为舍入误差 [3]。
用一段代码来展示：
>>># FP32相加不会有问题。 >>>torch.tensor(2**-3) torch.tensor(2**-14) tensor(0.1251) >>># FP相加，较小的数字会被忽略。因为在[2**-3, 2**-2]内，FP16表示固定间隔为2**-13。 >>>#也就是说比2**-下一个下一个数字是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)如果我们在训练过程中 FP32 替代为 FP有以下两个好处：
1. 减少显存占用：FP16 只占用显存 FP32 这使我们能够使用更大的一半 batch size；
2. 加速训练：使用 FP16.模型的训练速度几乎可以提高 1 倍。
如果我们简单地从模型权重和输入 FP32 转化成 FP16.虽然速度可以翻倍，但模型的精度会受到严重影响。原因如下:
FP16 表示范围不大，超过 数字会溢出变成 inf，小于 数字会溢出变成 0.溢出更为常见，因为在网络训练的后期，模型的梯度往往很小，甚至小于 FP16 的下限 ，梯度值会变成 0.模型参数无法更新。 SSD 训练过程中网络的梯度统计有 67% 溢出值变成 0。
即使梯度不会上下溢出，如果梯度值与模型参数值相差太远，也会出现舍入误差。假设模型参数 weight ，学习率 ，梯度 gradient ，weight weight gradient 。
为了解决溢出问题，论文中计算出来了 loss 值进行缩放 (scale)，由于链法的存在，对 loss 缩放将作用于每个梯度。缩放后的梯度将平移到 FP16 有效范围内。这样就可以用了。 FP16 存储梯度不会溢出。另外，更新前需要
注意这里一定要先转成 FP32，不然 unscale 它仍然会溢出。
缩放因子 (loss_scale) 框架通常是自动确定的，只要没有发生 inf 或者 nan，loss_scale 越大越好。因为随着训练的进行，网络的梯度会越来越小，越大 loss_scale 可更充分利用 FP16 表示范围。
为了实现 FP16 在训练中，我们需要将模型权重和输入数据转换为 FP16.反向传播会得到 FP16 梯度。如果此时直接更新，因为值往往很小，与模型权重的差距会很大，可能会出现放弃误差的问题。
所以解决办法是：将等数据用 同时和维护一份 的用于更新。反向传播。 FP16 梯度后，，最后更新 FP32 模型权重。因为整个更新过程是在 FP32 在环境中进行，因此不会出现舍入误差。
对于那些在 FP16 环境中运行不稳定的模块，我们会将其添加到黑名单中，强制它在 FP32 的精度下运行。比如需要计算 batch 均值的 BN 层就应该在 FP32 下运行，否则会发生舍入误差。还有一些函数对于算法精度要求很高，比如 torch.acos()，也应该在 FP32 下运行。论文中的黑名单只包含 BN 层。
如何保证黑名单模块在 FP32 环境中运行：以 BN 层为例，将其权重转为 FP32，并且将输入从 FP16 转成 FP32，这样就可以保证整个模块是在 FP32 下运行的。
Tensor Core 可以让 FP16 做矩阵相乘，然后把结果累加到 FP32 的矩阵中。这样既可以享受 FP16 高速的矩阵乘法，又可以利用 FP32 来消除舍入误差。
搞不懂 Tensor Core 是如何应用到 AMP 中的。有人说 Tensor Core 可以帮助我们利用 FP16 的梯度来更新 FP32 的模型权重。但是阅读了 apex 的源码之后，我发现 FP16 的梯度会先转化为 FP32，再做更新，所以权重更新和 Tensor Core 并无关系。以后弄明白了再回来补充吧。
其实将 FP16 和 FP32 混合起来使用是必然的结果，有以下几个原因：
1. 在网络训练的后期，梯度的值非常小，可能会让 FP16 下溢出。如果不使用 FP32，即使我们通过缩放操作暂时规避了这个问题，权重更新时的 unscale 操作还是会让梯度下溢出；
2. 承接第 1 条，就算梯度能够以 FP16 表示，但是可能会下溢出。所以权重更新这步操作还是得在 FP32 下运行；
3. 承接第 2 条，就算不会下溢出，其值相对于权重本身也是非常小的。这步操作中可能会发生舍入误差的问题；
4. 承接第 3 条，就算不会发生舍入误差，有些算子在 FP16 下也是不稳定的，比如 BN、torch.acos 等。
首先介绍下 apex 提供的几种 opt-level: o1, o2, o3, o4。注意这里是不是数字"0"。
图片来自：全网最全-混合精度训练原理
https://zhuanlan.zhihu.com/p/441591808
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。
1. 根据黑白名单对 PyTorch 内置的函数进行包装 [4]。白名单函数强制 FP16，黑名单函数强制 FP32。其余函数则根据参数类型自动判断，如果参数都是 FP16，则以 FP16 运行，如果有一个参数为 FP32，则以 FP32 运行。
2. 将 loss_scale 初始化为一个很大的值 [5]。
3. 对于每次迭代
(a). 前向传播：模型权重是 FP32，按照黑白名单自动选择算子精度。
(b). 将 loss 乘以 loss_scale [6]
(c). 反向传播，因为模型权重是 FP32，所以即使函数以 FP16 运行，也会得到 FP32 的梯度。
(d). 将梯度 unscale [7]，即除以 loss_scale
(e). 如果检测到 inf 或 nan [8]
i. loss_scale /= 2 [9]
ii. 跳过此次更新 [10]
(f). optimizer.step()，执行此次更新
(g). 如果连续2000次迭代都没有出现 inf 或 nan，则 loss_scale *= 2 [11]
1. 将除了 BN 层以外的模型权重转化为 FP16 [12]，并且包装了 forward 函数 [13]，将其参数也转化为 FP16；
2. 维护一个 FP32 的模型权重副本用于更新 [14]；
3. 将 loss_scale 初始化为一个很大的值 [15]；
4. 对于每次迭代
(a). 前向传播: 除了 BN 层是 FP32，模型其它部分都是 FP16。
(b). 将 loss 乘以 loss_scale [16]
(c). 反向传播，得到 FP16 的梯度
(d). 将 FP16 梯度转化为 FP32，并 unscale [17]
(e). 如果检测到 inf 或 nan [18]
i. loss_scale /= 2 [19]
ii. 跳过此次更新 [20]
(f). optimizer.step()，执行此次更新
(g). 如果连续2000次迭代都没有出现 inf 或 nan，则 loss_scale *= 2 [21]
此外，还推荐阅读 MMCV 对于 AMP 的 o2 实现 [22]，代码比 apex 更加清晰。但因为我想同时讲 o1 和 o2，就没有选择解读 MMCV 的代码，有兴趣的读者可以进一步研究。
[1] https://arxiv.org/abs/1710.03740
[2] https://github.com/NVIDIA/apex
[3] https://en.wikipedia.org/wiki/Round-off_error#Addition
[4] https://github.com/NVIDIA/apex/blob/1403c21acf87b0f2245278309071aef17d80c13b/apex/amp/amp.py#L68
[5] https://github.com/NVIDIA/apex/blob/1403c21acf87b0f2245278309071aef17d80c13b/apex/amp/scaler.py#L40
[6] https://github.com/NVIDIA/apex/blob/1403c21acf87b0f2245278309071aef17d80c13b/apex/amp/handle.py#L113
[7] https://github.com/NVIDIA/apex/blob/1403c21acf87b0f2245278309071aef17d80c13b/apex/amp/_process_optimizer.py#L123
[8] https://github.com/NVIDIA/apex/blob/1403c21acf87b0f2245278309071aef17d80c13b/apex/amp/scaler.py#L202
[9] https://github.com/NVIDIA/apex/blob/1403c21acf87b0f2245278309071aef17d80c13b/apex/amp/scaler.py#L207
[10] https://github.com/NVIDIA/apex/blob/1403c21acf87b0f2245278309071aef17d80c13b/apex/amp/handle.py#L128
[11] https://github.com/NVIDIA/apex/blob/1403c21acf87b0f2245278309071aef17d80c13b/apex/amp/scaler.py#L213
[12] https://github.com/NVIDIA/apex/blob/1403c21acf87b0f2245278309071aef17d80c13b/apex/amp/_initialize.py#L179
[13] https://github.com/NVIDIA/apex/blob/1403c21acf87b0f2245278309071aef17d80c13b/apex/amp/_initialize.py#L194
[14] https://github.com/NVIDIA/apex/blob/1403c21acf87b0f2245278309071aef17d80c13b/apex/amp/_process_optimizer.py#L44
[15] https://github.com/NVIDIA/apex/blob/1403c21acf87b0f2245278309071aef17d80c13b/apex/amp/scaler.py#L40
[16] https://github.com/NVIDIA/apex/blob/1403c21acf87b0f2245278309071aef17d80c13b/apex/amp/handle.py#L113
[17] https://github.com/NVIDIA/apex/blob/1403c21acf87b0f2245278309071aef17d80c13b/apex/amp/scaler.py#L94
[18] https://github.com/NVIDIA/apex/blob/1403c21acf87b0f2245278309071aef17d80c13b/apex/amp/scaler.py#L202
[19] https://github.com/NVIDIA/apex/blob/1403c21acf87b0f2245278309071aef17d80c13b/apex/amp/scaler.py#L207
[20] https://github.com/NVIDIA/apex/blob/1403c21acf87b0f2245278309071aef17d80c13b/apex/amp/handle.py#L128
[21] https://github.com/NVIDIA/apex/blob/1403c21acf87b0f2245278309071aef17d80c13b/apex/amp/scaler.py#L213
[22] https://github.com/open-mmlab/mmcv/blob/f5425ab7611ab2376ddb478b57cb2f46f6054e13/mmcv/runner/hooks/optimizer.py#L344
[23] https://docs.nvidia.com/deeplearning/performance/mixed-precision-training/index.html
[24] https://pytorch.org/docs/stable/amp.html#autocast-op-reference
[25] PyTorch必备神器 | 唯快不破：基于Apex的混合精度加速
[26] https://zhuanlan.zhihu.com/p/103685761
[27] https://zhuanlan.zhihu.com/p/441591808
如何才能让更多的优质内容以更短路径到达读者群体，缩短读者寻找优质内容的成本呢？
总有一些你不认识的人，知道你想知道的东西。PaperWeekly 或许可以成为一座桥梁，促使不同背景、不同方向的学者和学术灵感相互碰撞，迸发出更多的可能性。
PaperWeekly 鼓励高校实验室或个人，在我们的平台上分享各类优质内容，可以是，也可以是、或等。我们的目的只有一个，让知识真正流动起来。
📝
• 文章确系个人，未曾在公开渠道发表，如为其他平台已发表或待发表的文章，请明确标注
• 稿件建议以 格式撰写，文中配图以附件形式发送，要求图片清晰，无版权问题
• PaperWeekly 尊重原作者署名权，并将为每篇被采纳的原创首发稿件，提供，具体依据文章阅读量和文章质量阶梯制结算
📬
• 投稿邮箱：hr@paperweekly.site
• 来稿请备注即时联系方式（微信），以便我们在稿件选用的第一时间联系作者
• 您也可以直接添加小编微信（）快速投稿，备注：姓名-投稿
🔍
现在，在也能找到我们了
进入知乎首页搜索
点击订阅我们的专栏吧
·
·
·