【训练技巧】torch.cuda.amp.GradScaler() 深入详解

最新推荐文章于 2025-09-26 10:44:10 发布

原创最新推荐文章于 2025-09-26 10:44:10 发布 · 937 阅读

19 ·

CC 4.0 BY-SA版权

文章标签：

#人工智能 #计算机视觉 #机器学习 #算法

模型训练及技巧类专栏收录该内容

21 篇文章

订阅专栏

`torch.cuda.amp.GradScaler()` 深入详解

1. 核心作用

GradScaler 是 PyTorch 自动混合精度（Automatic Mixed Precision, AMP）训练的核心组件，主要解决 float16 数值精度不足 的问题：

float16 的表示范围（ $\times 10^{-5} \sim 65504$ ）远小于 float32（ $\times 10^{-45} \sim 3 \times 10^{38}$ ）
当梯度值 $\times 10^{-5}$ 时，float16 会将其视为 0（下溢），导致权重无法更新
当梯度值 $g > 65504$ 时，float16 会溢出为 inf（上溢），破坏训练过程

GradScaler 通过 动态缩放梯度 将梯度值保持在 float16 的安全范围内：
$g_{\text{scaled}} = s \cdot g$
其中 $s$ 是缩放因子（scale factor）， $g$ 是原始梯度。

2. 工作原理

(1) 梯度缩放

在反向传播前对损失函数进行缩放：

scaled_loss = scaler.scale(loss)  # loss -> s * loss
scaled_loss.backward()            # 梯度 = s * ∇loss

此时梯度被放大 $s$ 倍，避免了下溢风险。

(2) 梯度反缩放

在优化器更新前：

scaler.step(optimizer)  # 1. 梯度反缩放: g = g_scaled / s
                        # 2. 执行 optimizer.step()

梯度恢复原始量级： $\frac{g_{\text{scaled}}}{s}$
使用 float32 精度更新权重（避免精度损失）

(3) 缩放因子动态调整

scaler.update()  # 根据梯度状态调整 s

增大 $s$ ：若连续 $N$ 次未出现 inf/NaN（默认 $N = 2000$ ）
减小 $s$ ：若检测到 inf/NaN 梯度（通常减半）
初始 $s$ 默认为 $2^{16}$ （65536）

3. 数学意义

设损失函数为 $L(θ)\mathcal{L}(\theta)$ ，优化过程为：
$\theta_{t+1} = \theta_t - \eta \cdot \nabla\mathcal{L}(\theta_t)$
引入缩放后：
$\theta_{t+1} = \theta_t - \eta \cdot \frac{1}{s} \nabla(s \cdot \mathcal{L}(\theta_t))$
由于标量乘法与梯度线性兼容：
$\nabla(s \cdot \mathcal{L}) = s \cdot \nabla\mathcal{L}$
因此更新公式等价于：
$\theta_{t+1} = \theta_t - \eta \cdot \nabla\mathcal{L}(\theta_t)$
缩放操作不影响优化方向，仅避免数值问题。

4. 使用示例

from torch.cuda.amp import GradScaler, autocast

scaler = GradScaler()  # 初始化缩放器

for data, target in dataloader:
    optimizer.zero_grad()
    
    with autocast():  # 自动混合精度上下文
        output = model(data)  # float16 计算
        loss = loss_fn(output, target)
    
    # 缩放梯度 + 反向传播
    scaler.scale(loss).backward()
    
    # 反缩放 + 更新权重
    scaler.step(optimizer)
    
    # 调整缩放因子
    scaler.update()