PyTorch的自动混合精度amp

 参考自

PyTorch 源码解读之 torch.cuda.amp: 自动混合精度详解 - 知乎

 加入了我自己的一点理解，供自己复习时观看

自动混合精度原理

Tensor的dtype类型会自动变化，也就是框架按需自动调整cuda tensor的dtype

自动混合精度作用

低精度优势在于存储小、计算也会变快。

torch.cuda.amp

开启自动混合精度

torch.cuda.amp.GradScaler

通过放大loss的值来防止梯度的underflow下溢出（这只是BP的时候传递梯度信息使用，真正更新权重的时候还是要把放大的梯度再unscale回去）；

使用方法

from torch.cuda.amp import autocast as autocast

# 创建model，默认是torch.FloatTensor
model = Net().cuda()
optimizer = optim.SGD(model.parameters(), ...)

# 在训练最开始之前实例化一个GradScaler对象
scaler = GradScaler()

for epoch in epochs:
    for input, target in data:
        optimizer.zero_grad()

        # 前向过程(model + loss)开启 autocast
        with autocast():
            output = model(input)
            loss = loss_fn(output, target)

        # Scales loss. 为了梯度放大，扩大后得到的loss再执行反向传播
        scaler.scale(loss).backward()

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

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

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

1，如果出现了inf或者NaN，scaler.step(optimizer)会忽略此次的权重更新（optimizer.step() )，并且将scaler的大小缩小（乘上backoff_factor）；

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

也就是1.正向传播要开启混合精度，2.计算loss时要先放缩loss， 3.使用scaler.step(optimizer)代替optimizer.step进行权重更新  4.根据3的结果看看scaler是否要更新，更新scaler

为什么autocast中只有向前传播？

autocast上下文应该只包含网络的前向过程（包括loss的计算），而不要包含反向传播，因为BP的op会使用和前向op相同的类型。