Apex：PyTorch混合精度训练

from apex import amp

model, optimizer = amp.initialize(model, optimizer, opt_level= "O1") 
# 这里是“欧一”，不是“零一”

with amp.scale_loss(loss, optimizer) as scaled_loss:
	scaled_loss.backward()

Notice: PyTorch >= 1.6

from torch.cuda.amp import autocast as autocast, GradScaler
 
# 在训练最开始之前实例化一个GradScaler对象
scaler = GradScaler()

optimizer.zero_grad()

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

# Scales loss，这是因为半精度的数值范围有限，因此需要用它放大
scaler.scale(loss).backward()

# scaler.step() unscale之前放大后的梯度，但是scale太多可能出现inf或NaN
# 故其会判断是否出现了inf/NaN
# 如果梯度的值不是 infs 或者 NaNs, 那么调用optimizer.step()来更新权重,
# 如果检测到出现了inf或者NaN，就跳过这次梯度更新，同时动态调整scaler的大小
scaler.step(optimizer)

# 查看是否要更新scaler,这个要注意不能丢
scaler.update()