Pytorch中的自动混合精度

最近在训练yolo v5的模型时，出现了这样一个bug：

 cannot import name 'amp'

这个模块之前从来没有见过，所以就去了解了一下，发现是pytorch中的自动混合精度模块。这是yolov5新使用的技术，v4，v3都没有出现过。

1.什么是自动混合精度

自动混合精度（automatic mixed precision (AMP)）是在pytorch1.6版本中发布的。

神经网络计算框架的核心就是Tensor， 在深度学习中，Tensor实际上就是一个多维数组（multidimensional array），其目的是能够创造更高维度的矩阵、向量。

Tensor有不同的数据类型，在pytorch中，Tensor有十种类型：

torch.FloatTensor (32-bit floating point)
torch.DoubleTensor (64-bit floating point)
torch.HalfTensor (16-bit floating point 1)
torch.BFloat16Tensor (16-bit floating point 2)
torch.ByteTensor (8-bit integer (unsigned))
torch.CharTensor (8-bit integer (signed))
torch.ShortTensor (16-bit integer (signed))
torch.IntTensor (32-bit integer (signed))
torch.LongTensor (64-bit integer (signed))
torch.BoolTensor (Boolean)

我们创建的Tensor默认的类型为32-bit floating point，这就是32位浮点型精度的Tensor。

在自动混合精度中，我们主要关注两种类型的Tensor，他们分别是torch.FloatTensor和torch.HalfTensor，即混合精度。

自动混合精度中的“自动”表示Tensor的类型，即dtype会自动变化，框架会按照需要自己调整tensor的dtype。当然我们也可以手动调整。

导入amp模块:from torch.cuda import amp

torch.cuda.amp 的名字意味着这个功能只能在cuda上使用，事实上，这个功能正是NVIDIA的开发人员贡献到PyTorch项目中的。而只有支持Tensor core的CUDA硬件才能享受到AMP的好处（比如2080ti显卡）。Tensor Core是一种矩阵乘累加的计算单元，每个Tensor Core每个时钟执行64个浮点混合精度操作（FP16矩阵相乘和FP32累加），英伟达宣称使用Tensor Core进行矩阵运算可以轻易的提速，同时降低一半的显存访问和存储。

因此，在PyTorch中，当我们提到自动混合精度训练，我们说的就是在NVIDIA的支持Tensor core的CUDA设备上使用torch.cuda.amp.autocast （以及torch.cuda.amp.GradScaler）来进行训练。

2.自动混合精度出现的意义

在训练的时候为什么要在torch.FloatTensor类型的基础上，混合另一种数据类型（torch.HalfTensor）呢？

那就是在某些情况下，我们用torch.HalfTensor会比torch.FloatTensor有优势。

torch.HalfTensor的优势就是存储小、计算快、更好的利用CUDA设备的Tensor Core。因此训练的时候可以减少显存的占用（可以增加batchsize了），同时训练速度更快；

torch.HalfTensor的劣势就是：数值范围小（更容易Overflow / Underflow）、舍入误差（Rounding Error，导致一些微小的梯度信息达不到16bit精度的最低分辨率，从而丢失）。

为了消除torch.HalfTensor的劣势，又引入了一个新torch.cuda.amp.GradScaler，它是通过放大loss的值来防止梯度的underflow（这只是BP的时候传递梯度信息使用，真正更新权重的时候还是要把放大的梯度再unscale回去）

3.自动混合精度训练的使用

我们看看自动混合精度训练是在yolov5中怎么使用的：

在训练最开始之前实例化一个GradScaler对象。

这是一个autocast的上下文管理器，是为了在前向传播时，使用自动混合精度，当进入autocast的上下文后，上面列出来的那些CUDA ops 会把tensor的dtype转换为半精度浮点型，从而在不损失训练精度的情况下加快运算。刚进入autocast的上下文时，tensor可以是任何类型，你不要在model或者input上手工调用.half() ，框架会自动做，这也是自动混合精度中“自动”一词的由来。

autocast上下文只用包含前向传播的过程，因为在反向传播时，会使用和前向传播时一样的数据类型。

在反向传播时，就会用到之前我们实例化的GradScaler对象。

        #Scales loss. 为了梯度放大.
        scaler.scale(loss).backward()

        # scaler.step() 首先把梯度的值unscale回来.
        # 如果梯度的值不是 infs 或者 NaNs, 那么调用optimizer.step()来更新权重,
        # 否则，忽略step调用，从而保证权重不更新（不被破坏）
        scaler.step(optimizer)

        # 准备着，看是否要增大scaler
        scaler.update()

4.补充

如下操作中tensor会被自动转化为半精度浮点型的torch.HalfTensor：

matmul
addbmm
addmm
addmv
addr
baddbmm
bmm
chain_matmul
conv1d
conv2d
conv3d
conv_transpose1d
conv_transpose2d
conv_transpose3d
linear
matmul
mm
mv
prelu

其他操作需要手动转换为自动混合精度。