混合精度训练 AMP 实践
混合精度训练（AMP）通过在 FP16 和 FP32 之间自动切换，在不损失模型质量的前提下显著加速训练并降低内存占用。
AMP 使用
python
from torch.cuda.amp import autocast, GradScaler
scaler = GradScaler()
model = MyModel().cuda()
optimizer = torch.optim.AdamW(model.parameters())
for batch in dataloader:
optimizer.zero_grad()
with autocast():
output = model(batch)
loss = criterion(output, target)
scaler.scale(loss).backward()
scaler.unscale_(optimizer)
torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
scaler.step(optimizer)
scaler.update()GradScaler 的作用
FP16 的表示范围有限，容易导致梯度下溢（变为零）。GradScaler 通过放大损失来避免下溢，再在更新前缩小梯度。