torch.cuda.amp.autocast()使用示例

# 定义模型和优化器
model = Net().cuda()
optimizer = optim.SGD(model.parameters(), ...)

# 在训练最开始定义GradScalar的实例
scaler = GradScaler()

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

        # 利用with语句，在autocast实例的上下文范围内，进行模型的前向推理和loss计算
        with autocast():
            output = model(input)
            loss = loss_fn(output, target)

        # 对loss进行缩放，针对缩放后的loss进行反向传播
        # （此部分计算在autocast()作用范围以外）
        scaler.scale(loss).backward()

        # 将梯度值缩放回原尺度后，优化器进行一步优化
        scaler.step(optimizer)

        # 更新scalar的缩放信息
        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）。