混合精度训练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的组成分为了三个部分：

最高位表示符号位；
有5位表示exponent位;
有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 - 博客园