1. 自动混合精度理论概要

一句话概括的是，自动混合精度的实现使用的autocast + GradScaler。以下是对自动混合精度的介绍：

1. amp：Automatic mixed precision，自动混合精度，可以在神经网络推理过程中，针对不同的层，采用不同的数据精度进行计算，从而实现节省显存和加快速度的目的。自动混合精度的关键词有两个：自动、混合精度。这是由PyTorch 1.6的torch.cuda.amp模块带来的：from torch.cuda import amp
2. 混合精度：预示着有不止一种精度的Tensor，那在PyTorch的AMP模块里是几种呢？2种：torch.FloatTensor（浮点型 32位）和torch.HalfTensor（半精度浮点型 16位）；
3. 自动：预示着Tensor的dtype类型会自动变化，也就是框架按需自动调整tensor的dtype（其实不是完全自动，有些地方还是需要手工干预）；

补充说明：

• torch.cuda.amp 的名字意味着这个功能只能在cuda上使用；
• torch默认的tensor精度类型是torch.FloatTensor

问题：为什么需要自动混合精度，也就是torch.FloatTensor和torch.HalfTensor的混合，而不全是torch.FloatTensor？或者全是torch.HalfTensor？

原因： 在某些上下文中torch.FloatTensor有优势，在某些上下文中torch.HalfTensor有优势。什么是torch.HalfTensor？

torch.HalfTensor：

• torch.HalfTensor的优势就是存储小、计算快、更好的利用CUDA设备的Tensor Core。因此训练的时候可以减少显存的占用（可以增加batchsize了），同时训练速度更快；
• torch.HalfTensor的劣势就是：数值范围小（更容易Overflow / Underflow）、舍入误差（Rounding Error，导致一些微小的梯度信息达不到16bit精度的最低分辨率，

     从而丢失）。

可见，当有优势的时候就用torch.HalfTensor，而为了消除torch.HalfTensor的劣势，有两种解决方案：

• 方案一：

梯度scale，这正是上一小节中提到的torch.cuda.amp.GradScaler，通过放大loss的值来防止梯度消失underflow（这只是BP的时候传递梯度信息使用，真正更新权重的时候还是要把放大的梯度再unscale回去）；

• 方案二：

回落到torch.FloatTensor，这就是混合一词的由来。那怎么知道什么时候用torch.FloatTensor，什么时候用半精度浮点型呢？这是PyTorch框架决定的，AMP上下文中，一些常用的操作中tensor会被自动转化为半精度浮点型的torch.HalfTensor（如：conv1d、conv2d、conv3d、linear、prelu等）。实际上，在训练过程中，内存中占据大部分的基本都是 activations 的值。特别是在batchsize 很大的情况下， activations 更是特别占据空间。 保存 activiations 主要是为了在 back-propogation 的时候进行计算。因此，只要 activation 的值基本都是使用 fp16 来进行存储的话，则最终模型与 fp32 相比起来， 内存占用也基本能够减半。

总结：

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

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 之间皆可。

自动混合精度的实现使用的autocast + GradScaler，下面分别对autocast和GradScaler进行介绍。

2. Autocast使用介绍

使用torch.cuda.amp模块中的autocast 类

Autocast使用的参考代码：

from torch.cuda import amp
# 创建model，默认是torch.FloatTensor
model = Net().cuda()
optimizer = optim.SGD(model.parameters(), ...)
# 判断能否使用自动混合精度
enable_amp = True if "cuda" in device.type else False
for input, target in data:
    optimizer.zero_grad()
    # 前向过程(model + loss)开启 autocast
    with amp.autocast(enabled=enable_amp):
        output = model(input)
        loss = loss_fn(output, target)
    # 反向传播在autocast上下文之外
    loss.backward()
    optimizer.step()

需要注意:

• 当进入autocast，自动将torch.FloatTensor类型转化为torch.HalfTensor，而不需要手动设置model.half()/input.half，框架会自动做，这也是自动混合精度中“自动”一词的由来。
• autocast上下文应该只包含网络的前向过程（包括loss的计算），而不要包含反向传播。

3. GradScaler使用介绍

这里GradScaler就是第二小节中提到的梯度scaler模块，需要在训练最开始之前使用amp.GradScaler实例化一个GradScaler对象

GradScaler使用的参考代码：

from torch.cuda import amp
# 创建model，默认是torch.FloatTensor
model = Net().cuda()
optimizer = optim.SGD(model.parameters(), ...)
# 判断能否使用自动混合精度
enable_amp = True if "cuda" in device.type else False
# 在训练最开始之前实例化一个GradScaler对象
scaler = amp.GradScaler(enabled=enable_amp)
for epoch in epochs:
    for input, target in data:
        optimizer.zero_grad()
        # 前向过程(model + loss)开启 autocast
        with amp.autocast(enabled=enable_amp):
            output = model(input)
            loss = loss_fn(output, target)
        # 1、Scales loss.  先将梯度放大 防止梯度消失
        scaler.scale(loss).backward()
        # 2、scaler.step()   再把梯度的值unscale回来.
        # 如果梯度的值不是 infs 或者 NaNs, 那么调用optimizer.step()来更新权重,
        # 否则，忽略step调用，从而保证权重不更新（不被破坏）
        scaler.step(optimizer)
        # 3、准备着，看是否要增大scaler
        scaler.update()
        # 正常更新权重
        optimizer.zero_grad()

需要注意：

scaler的大小在每次迭代中动态的估计，为了尽可能的减少梯度underflow，scaler应该更大；但是如果太大的话，半精度浮点型的tensor又容易overflow（变成inf或者NaN）。所以动态估计的原理就是在不出现inf或者NaN梯度值的情况下尽可能的增大scaler的值——在每次scaler.step(optimizer)中，都会检查是否又inf或NaN的梯度出现：

• 如果出现了inf或者NaN，scaler.step(optimizer)会忽略此次的权重更新（optimizer.step() )，并且将scaler的大小缩小（乘上backoff_factor）；
• 如果没有出现inf或者NaN，那么权重正常更新，并且当连续多次（growth_interval指定）没有出现inf或者NaN，则scaler.update()会将scaler的大小增加（乘上growth_factor）。

4. 自动混合精度实现代码

YOLOv3-SPP中的代码：

scaler = torch.cuda.amp.GradScaler() if opt.amp else None
def train_one_epoch(model, optimizer, data_loader, device, epoch,
                    print_freq, accumulate, img_size,
                    grid_min, grid_max, gs,
                    multi_scale=False, warmup=False, scaler=None):
    ...
    # 使用自动混合精度
    with amp.autocast(enabled=scaler is not None):
        pred = model(imgs)
        # loss计算
        loss_dict = compute_loss(pred, targets, model)
        losses = sum(loss for loss in loss_dict.values())
    # backward
    # 通过放大loss的值来防止梯度的underflow，只在BP时传递梯度信息使用
    if scaler is not None:
        scaler.scale(losses).backward()     # 先将梯度放大 防止梯度消失
    else:
        losses.backward()
    # optimize
    # 每训练64张图片更新一次权重
    if ni % accumulate == 0:
        if scaler is not None:
            scaler.step(optimizer)  # 把梯度的值unscale回来
            scaler.update()
        else:
            optimizer.step()      
        # 正常更新权重
        optimizer.zero_grad()
    ...

主要原理：

Loss Scale 主要是为了解决 fp16 underflow 的问题。刚才提到，训练到了后期，梯度（特别是激活函数平滑段的梯度）会特别小，fp16 表示容易产生 underflow 现象。

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

多GPU训练：

单卡训练的话上面的代码已经够了。要是想多卡跑的话仅仅这样还不够，会发现在forward里面的每个结果都还是float32的。

解决办法：只要把model中的forward里面的代码用autocast代码块方式运行就好了

class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()
    def forward(self, input):
        with autocast():
            .....
        return

参考资料：

1. torch.cuda.amp自动混合精度训练 —— 节省显存并加快推理速度

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