2023年11月26日发(作者：)

Pytorch⾃动混合精度（AMP）训练

相关问题：

从1.6版本开始，已经内置了，采⽤⾃动混合精度训练就不需要加载第三⽅NVIDIA的apex库了。AMP -- （automatic

mixed-precision training）

⼀ 什么是⾃动混合训练(AMP)

默认情况下，⼤多数框架都采⽤32位浮点算法进⾏训练。2017年，NVIDIA研究了⼀种⽤于混合精度训练的⽅法，该⽅法在训练⽹络时将

单精度（FP32）与半精度(FP16)结合在⼀起，并使⽤相同的超参数实现了与FP32⼏乎相同的精度。

1.1、FP16理论基础：

在介绍AMP之前，先来理解下FP16与FP32，FP16也即半精度是⼀种计算机使⽤的⼆进制数据类型，使⽤2字节存储。⽽FLOAT就是

FP32。

在上图可以看到，与单精度float（32bit，4个字节）相⽐，半进度float16仅有16bit，2个字节组成。天然的存储空间是float的⼀半。 其

中，float16的组成分为了三个部分：

1. 最⾼位表⽰符号位，sign 位表⽰正负

2. 有5位表⽰exponent位， exponent 位表⽰指数

FP16的表⽰范例

总结：

FP16值动态区间

FP32值动态区间：

float16 最⼤范围是 [-65504，66504]，float16 能表⽰的精度范围是，超过这个数值的数字会被直接置0；

Tensor

默认Tensor是32bit floating point，这就是32位浮点型精度的tensor。

AMP(⾃动混合精度）的关键词有两个：⾃动，混合精度。这是由PyTorch 1.6的.模块带来的：

⾃动：Tensor的dtype类型会⾃动变化，框架按需⾃动调整tensor的dtype，当然有些地⽅还需⼿动⼲预。

混合精度：采⽤不⽌⼀种精度的Tensor，ensor和nsor

pytorch1.6的新包：， 的名字意味着这个功能只能在cuda上使⽤，是ＮVIDIA开发⼈员贡献到pytorch

⾥的。只有⽀持tensor core的CUDA硬件才能享受到AMP带来的优势（⽐如2080ti显卡）。

Tensor core是⼀种矩阵乘累加的计算单元，每个tensor core时针执⾏64个浮点混合精度操作（FP16矩阵相乘和FP32累加）。英伟达

宣称使⽤Tensor Core进⾏矩阵运算可以轻易的提速，同时降低⼀半的显存访问和存储。

因此，在PyTorch中，当我们提到⾃动混合精度训练，我们说的就是在NVIDIA的⽀持Tensor core的CUDA设备上使⽤

下溢出

2、舍⼊误差

舍⼊误差指的是当梯度过⼩时，⼩于当前区间内的最⼩间隔时，该次梯度更新可能会失败。

1.3 解决问题的办法：混合精度训练+动态损失放⼤

1、混合精度训练

在某些模型中，fp16矩阵乘法的过程中，需要利⽤ fp32 来进⾏矩阵乘法中间的累加(accumulated)，然后再将 fp32 的值转化为 fp16

进⾏存储。 换句不太严谨的话来说，也就是在内存中⽤FP16做储存和乘法从⽽加速计算，⽽⽤FP32做累加避免舍⼊误差。混合精度训练

的策略有效地缓解了舍⼊误差的问题。

在这⾥也就引出了，为什么⽹上⼤家都说，只有 Nvidia Volta 结构的 拥有 TensorCore 的CPU(例如V100)，才能利⽤ fp16 混合精度来

进⾏加速。 那是因为 TensorCore 能够保证 fp16 的矩阵相乘，利⽤ fp16 or fp32 来进⾏累加。在累加阶段能够使⽤ FP32 ⼤幅减少

混合精度训练的精度损失。⽽其他的GPU 只能⽀持 fp16 的 multiply-add operation。这⾥直接贴出原⽂句⼦：

Whereas previous GPUs supported only FP16 multiply-add operation, NVIDIA Volta GPUs introduce Tensor Cores that

multiply FP16 input matrices andaccumulate products into either FP16 or FP32 outputs

可以看到，其他所有值（weights，activations， gradients）均使⽤ fp16 来存储，⽽唯独权重weights需要⽤ fp32 的格式额外备份⼀

次。 这主要是因为，在更新权重的时候，往往公式: 权重 = 旧权重 + lr * 梯度，⽽在深度模型中，lr * 梯度 这个值往往是⾮常⼩的，如果利

⽤ fp16 来进⾏相加的话， 则很可能会出现上⾯所说的『舍⼊误差』的这个问题，导致更新⽆效。因此上图中，通过将weights拷贝成

fp32 格式，并且确保整个更新（update）过程是在 fp32 格式下进⾏的。

看到这⾥，可能有⼈提出这种 fp32 拷贝weight的⽅式，那岂不是使得内存占⽤反⽽更⾼了呢？是的， fp32 额外拷贝⼀份 weight 的确新

增加了训练时候存储的占⽤。 但是实际上，在训练过程中，内存中占据⼤部分的基本都是 activations 的值。特别是在batchsize 很⼤的情

况下， activations 更是特别占据空间。 保存 activiations 主要是为了在 back-propogation 的时候进⾏计算。因此，只要 activation 的

值基本都是使⽤ fp16 来进⾏存储的话，则最终模型与 fp32 相⽐起来， 内存占⽤也基本能够减半。

2、损失放⼤（Loss scaling)

即使了混合精度训练，还是存在⽆法收敛的情况，原因是激活梯度的值太⼩，造成了下溢出（Underflow）。Loss Scale 主要是为了解决

不过，autocast上下⽂只能包含⽹络的前向过程(包括loss的计算），不能包含反向传播，因为BP的op会使⽤和前向op相同的类型。

当然，有时在autocast中的代码会报错：

1Traceback (most recent call last):

2......

3 File "/opt/conda/lib/python3.7/site-packages/torch/nn/modules/", line 722, in _call_impl

4 result = d(*input, **kwargs)

5......

6RuntimeError: expected scalar type float but found c10::Half

对于RuntimeError:expected scaler type float but found c10:Half,应该是个bug,可在tensor上⼿动调⽤.float()来让type匹配。

2.2 GradScaler

使⽤前，需要在训练最开始前实例化⼀个GradScaler对象，例程如下：

1from import autocast, GradScaler

2

3# 创建model，默认是ensor

4model = Net().cuda()

5optimizer = (ters(), ...)

6

7# 在训练最开始之前实例化⼀个GradScaler对象

8scaler = GradScaler()

9

10for epoch in epochs:

11 for input, target in data:

12 _grad()

13

14 # 前向过程(model + loss)开启 autocast

15 with autocast():

16 output = model(input)

17 loss = loss_fn(output, target)

18

19 # Scales loss. 为了梯度放⼤.