pytorch混合精度训练

1 混合精度

计算机中的浮点数表示，按照IEEE754可以分为三种，分别是半精度浮点数、单精度浮点数和双精度浮点数。三种格式的浮点数因占用的存储位数不同，能够表示的数据精度也不同。

Signed bit用于控制浮点数的正负，0表示正数，1表示负数；
Exponent部分用于控制浮点数的大小，以2为底进行指数运算；
Significand部分用于控制浮点数的精度，存储浮点数的有效数字。

从上图中可以看出，fp16、fp32和fp64具有不同的表示范围和数值精度，fp16数值范围和精度最低，存储空间也最小。fp64数值范围和精度最高，存储空间也最大。fp32居两者之间。

默认深度学习模型训练过程中都是使用fp32.

2 fp16训练的优劣

该部分内容引用自：https://zhuanlan.zhihu.com/p/79887894
使用fp16存储网络的权重值、激活值和梯度值进行网络训练，好处是：

1. 减少显存占用：上面的图已经很明显的可以看出，fp16的存储空间为fp32的一半，如果使用fp16进行训练，那么可以减少一半的显存占用，因此也就可以使用更大的batchsize进行大模型的训练；
2. 加快训练和推理速度：fp16可以提高模型的训练和推理的速度。

完全使用fp16进行训练的缺点：

1. 数值溢出：fp16的动态范围比fp32的动态范围小很多，因此在计算过程中很容易出现数值上溢出和下溢出的问题，溢出之后就会出现“Nan”的问题。在深度学习中，由于激活值的梯度往往比权重的梯度小，更容易出现数值下溢的情况。一旦出现了数值下溢，就会造成发生数值下溢的层及其前面层的权重无法更新，训练失败；
2. 舍入误差：舍入误差是指梯度过小，小于当前区间内的最小间隔时，该次梯度更新可能会失败，如下图所示：
   某次梯度更新值小于了当前权重区间的最小间隔时，更新无效。

3 混合精度训练的实现思路

混合精度训练的思路来自于nvidia和百度合作的论文：https://arxiv.org/pdf/1710.03740.pdf

下面部分的内容引用自：https://zhuanlan.zhihu.com/p/103685761

1. FP32权重备份
   实现思路为，网络前向计算过程中，权重、激活值和梯度都使用fp16进行存储，同时保留一份fp32格式的权重的副本用于进行参数更新。
   进行权重更新时，需要对权重的梯度乘以学习率，而学习率一般都是一个远小于1的数字，因此权重的更新值就更小。如果使用fp16进行权重更新，那么更容易引入舍入误差的问题造成参数无效更新。而这里将参数的更新量应用于fp32精度的权重上，因为fp32的数值精度远远大于fp16，也就避免了更新过程中的舍入误差问题。

   看到这里，可能有人提出这种 fp32 拷贝weight的方式，那岂不是使得内存占用反而更高了呢？是的， fp32 额外拷贝一份 weight 的确新增加了训练时候存储的占用。 但是实际上，在训练过程中，内存中占据大部分的基本都是 activations 的值。特别是在batchsize 很大的情况下， activations 更是特别占据空间。 保存 activiations 主要是为了在 back-propogation 的时候进行计算。因此，只要 activation 的值基本都是使用 fp16 来进行存储的话，则最终模型与 fp32 相比起来， 内存占用也基本能够减半。

2. Loss Scale
   Loss Scale 主要是为了解决 fp16 underflow 的问题。刚才提到，训练到了后期，梯度（特别是激活函数平滑段的梯度）会特别小，fp16 表示容易产生 underflow 现象。 下图展示了 SSD 模型在训练过程中，激活函数梯度的分布情况：可以看到，有67%的梯度小于$2^{-24}$，如果用 fp16 来表示，则这些梯度都会变成0。
   为了解决梯度过小的问题，论文中对计算出来的loss值进行scale，由于链式法则的存在，loss上的scale会作用也会作用在梯度上。这样比起对每个梯度进行scale更加划算。 scaled 过后的梯度，就会平移到 fp16 有效的展示范围内。

   这样，scaled-gradient 就可以一直使用 fp16 进行存储了。只有在进行更新的时候，才会将 scaled-gradient 转化为 fp32，同时将scale抹去。论文指出， scale 并非对于所有网络而言都是必须的。而scale的取值为也会特别大，论文给出在 8 - 32k 之间皆可。

3. 提高计算精度

   在论文中还提到一个『计算精度』的问题：在某些模型中，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 and accumulate products into either FP16 or FP32 outputs。

4 实现代码

4.1 nvidia版本

nvidia官方github：https://github.com/NVIDIA/apex

在pytorch中应用：

from apex import amp
model, optimizer = amp.initialize(model, optimizer, opt_level="O1") # 这里是“欧一”，不是“零一”
with amp.scale_loss(loss, optimizer) as scaled_loss:
    scaled_loss.backward()

参数 opt_level：

O0：纯FP32训练，可以作为accuracy的baseline；
O1：混合精度训练（推荐使用），根据黑白名单自动决定使用FP16（GEMM, 卷积）还是FP32（Softmax）进行计算；
O2：“几乎FP16”混合精度训练，不存在黑白名单，除了Batch norm，几乎都是用FP16计算；
O3：纯FP16训练，很不稳定，但是可以作为speed的baseline。

4.2 pytorch官方版本

pytorch1.6开始支持amp，位于torch.cuda.amp模块下，具体例子见https://pytorch.org/docs/stable/notes/amp_examples.html

包括 torch.cuda.amp.autocast和torch.cuda.amp.GradScaler两个模块，autocast针对选定的代码块自动选取适合的计算精度，以便在保持模型准确率的情况下最大化改善训练效率；GradScaler即进行梯度缩放，以最大程度避免使用fp16进行运算时的梯度下溢。

4.2.1 示例代码

# Creates model and optimizer in default precision
model = Net().cuda()
optimizer = optim.SGD