混合精度训练amp，torch.cuda.amp.autocast():

1 需要什么GPU：

在上面讲述了为什么利用混合精度加速，需要拥有 TensorCore 的GPU

0x02. 基础理论：

在日常中深度学习的系统，一般使用的是单精度 float（Single-Precision）浮点表示。在了解混合精度训练之前，我们需要先对其中的主角半精度『float16』进行一定的理论知识学习。

float vs float16 的组成bitmap

在上图可以看到，与单精度float（32bit，4个字节）相比，半进度float16仅有16bit，2个字节组成。天然的存储空间是float的一半。 其中，float16的组成分为了三个部分：

1. 最高位表示符号位；
2. 有5位表示exponent位;
3. 有10位表示fraction位;

在上面的例子中，我们可以对 float16 有一个感性的认识，对这个范围有一个大概的印象，后面会用到：

• float16 最大范围是 [-65504 - 66504]
• float16 能表示的精度范围是 2^-24 ，超过这个数值的数字会被直接置0；

为什么需要半精度：

float16和float相比恰里，总结下来就是两个原因：内存占用更少，计算更快。

• 内存占用更少：这个是显然可见的，通用的模型 fp16 占用的内存只需原来的一半。memory-bandwidth 减半所带来的好处：
  • 模型占用的内存更小，训练的时候可以用更大的batchsize。
  • 模型训练时，通信量（特别是多卡，或者多机多卡）大幅减少，大幅减少等待时间，加快数据的流通。
• 计算更快：
  • 目前的不少GPU都有针对 fp16 的计算进行优化。论文指出：在近期的GPU中，半精度的计算吞吐量可以是单精度的 2-8 倍；

为什么需要混合精度：

数据溢出问题：Overflow / Underflow

在第一节中，我们提到了 fp16 的有效的动态范围约为 (  [-65504 - 66504] )，比单精度的float要狭窄很多。对于深度学习而言，最大的问题在于 Underflow（下溢出），在训练后期，例如激活函数的梯度会非常小， 甚至在梯度乘以学习率后，值会更加小。（权重 = 旧权重 + lr * 梯度）

3.2.2 舍入误差（Rounding Error）

何为舍入误差，引用[2]中的一张图说的比较透彻：

这个例子非常直观的阐述了『舍入误差』这个说法。而至于上面提到的，FP16的最小间隔是一个比较玄乎的事，在wikipedia的引用上有这么一张图： 描述了 fp16 各个区间的最小gap。

半精度FP16不同区间的最小间隔

3.3 解决办法：

3.3.1 FP32 权重备份

这种方法主要是用于解决舍入误差的问题。其主要思路，可以概括为：weights, activations, gradients 等数据在训练中都利用FP16来存储，同时拷贝一份FP32的weights，用于更新。 在这里，我直接贴一张论文[1]的图片来阐述：

权重fp32备份策略

可以看到，其他所有值（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 相比起来， 内存占用也基本能够减半。

3.3.2 Loss Scale（损失放大）

即使了混合精度训练，还是存在无法收敛的情况，原因是激活梯度的值太小，造成了溢出。可以通过使用torch.cuda.amp.GradScaler，通过放大loss的值来防止梯度的underflow（只在BP时传递梯度信息使用，真正更新权重时还是要把放大的梯度再unscale回去）；

• 反向传播前，将损失变化手动增大2^k倍，因此反向传播时得到的中间变量（激活函数梯度）则不会溢出；
• 反向传播后，将权重梯度缩小2^k倍，恢复正常值

Loss Scale 主要是为了解决 fp16 underflow 的问题。刚才提到，训练到了后期，梯度（特别是激活函数平滑段的梯度）会特别小，fp16 表示容易产生 underflow 现象。 下图展示了 SSD 模型在训练过程中，激活函数梯度的分布情况：可以看到，有67%的梯度小于 -65504 ，如果用 fp16 来表示，则这些梯度都会变成0。

为了解决梯度过小的问题，论文中对计算出来的loss值进行scale，由于链式法则的存在，loss上的scale会作用也会作用在梯度上。这样比起对每个梯度进行scale更加划算。 scaled 过后的梯度，就会平移到 fp16 有效的展示范围内。

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

3.3.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。

浅谈混合精度训练 - 知乎

如何在pytorch中使用amp？

连接：Pytorch自动混合精度(AMP)介绍与使用 - jimchen1218 - 博客园