Training Technology Two : Torch.cuda.amp and Save

torch.cuda.amp是PyTorch中的一种自动混合精度训练工具。它可以帮助开发者在保持模型训练精度的同时，提高训练速度和减少内存占用。

# 导入必要的库
import torch
from torch.cuda.amp import autocast, GradScaler
 
# 创建GradScaler对象
scaler = GradScaler()
 
# 在训练循环中使用autocast和GradScaler
for data, target in train_loader:
    optimizer.zero_grad()
    
    # 开始自动混合精度训练
    with autocast():
        output = model(data)
        loss = loss_function(output, target)
    
    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()

在上述示例中，autocast()用于自动将前向传播、损失计算和反向传播的计算转换为半精度（FP16）运算。然后，GradScaler被用于缩放和反缩放梯度，以防止梯度下溢或溢出，并根据模型权重的尺度进行优化器步骤。

本文根据PyTorch 源码解读之 torch.cuda.amp: 自动混合精度详解, 节选部分知识.

autocast

可以作为 Python 上下文管理器和装饰器来使用，用来指定脚本中某个区域、或者某些函数，按照自动混合精度来运行。

GradScaler

fp16与fp32

fp16采用2字节存储, fp32采用4字节存储.

fp16第1位表示符号, 第2~6位表示指数, 第7~16为表示分数, 计算公式为:
$$(-1{)}^{signbit}\times 2^{(exponent-15)}\times (1+\frac{fraction}{1024})$$

详情参考该文

关于 unscale_

经过 scaler.scale(loss).backward()得到的梯度是 scaled gradient如果想要在 scaler.step(optimizer)前进行梯度裁剪等操作就必须先用 ``scaler.unscale_(optimizer) ```得到 unscaled gradient.

详情参考该文.

因此,一般有如下搭配:

scaler.scale(loss).backward()

# Unscales the gradients of optimizer's assigned params in-place
scaler.unscale_(optimizer)

# Since the gradients of optimizer's assigned params are unscaled, clips as usual:
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm)

# optimizer's gradients are already unscaled, so scaler.step does not unscale them,
# although it still skips optimizer.step() if the gradients contain infs or NaNs.
scaler.step(optimizer)

# Updates the scale for next iteration.
scaler.update()

保存

关于保存的过程可以参考该文

我们可以选择保存参数文件或者字典文件.

通常字典文件中包含的关键字有:

1. epoch;
2. state_dict;
3. optimizer;
4. 如果有scaler的话还要保存scaler.