PyTorch AMP实战：核心模块详解

好，咱们直接进入正题。PyTorch 的 AMP（自动混合精度）其实就三个核心组件：autocast 上下文管理器、GradScaler 梯度缩放器，以及它们配合起来的训练循环。我刚开始用 AMP 的时候，总觉得这玩意儿很玄乎，后来拆开一看，说白了就是「该用 fp16 的地方用 fp16，该用 fp32 的地方用 fp32」。

你想想看，模型里有些层对精度不敏感，比如卷积、全连接的计算，用半精度能快不少。但像 BatchNorm、Softmax 这些，精度一掉就崩。AMP 就是自动帮你做这个决策的。

torch.cuda.amp 模块概览

这个模块在 PyTorch 1.6 之后就成了官方标配。我个人习惯直接 from torch.cuda.amp import autocast, GradScaler，这样写起来干净。

它的核心逻辑其实就两步：

  • autocast：自动选择精度。前向传播时，它像个智能调度员，哪些操作用 fp16，哪些用 fp32，它说了算。
  • GradScaler：防止梯度下溢。fp16 的数值范围窄，小梯度容易变成 0。缩放器先把 loss 放大，反向传播完再缩回来。

嗯，这里要注意：autocast 只管前向和 loss 计算，反向传播它不管。梯度缩放才是专门处理反向传播那摊子事的。

autocast 上下文管理器

用法特别简单。你只需要把前向传播和 loss 计算包在 with autocast(): 里就行。我见过有人把整个训练循环都包进去，其实没必要——反向传播和参数更新不需要它。

from torch.cuda.amp import autocast

model = MyModel().cuda()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

for data, target in dataloader:
    data, target = data.cuda(), target.cuda()
    optimizer.zero_grad()
    
    # 只有这里需要 autocast
    with autocast():
        output = model(data)
        loss = loss_fn(output, target)
    
    loss.backward()
    optimizer.step()

为什么只包这一块？因为 autocast 的本质是重写某些 CUDA 操作的精度。你把它想成一个「精度切换开关」，只在你需要自动选择精度的代码段里打开它就行。

小技巧： 如果你用到了多个 GPU 或者分布式训练，autocast 依然有效。它只影响当前线程的 CUDA 操作，不会跨卡打架。

GradScaler 梯度缩放器

这个才是 AMP 的精华。fp16 的梯度范围大概在 2^-24 到 65504 之间，很多小梯度直接变成 0 了。我曾经在一个 NLP 项目里遇到过，模型训练到一半 loss 突然不降了，查了半天发现是梯度下溢——所有小梯度都被 fp16 吞掉了。

GradScaler 的做法很简单：

  1. 把 loss 乘上一个缩放因子（默认 2^16）
  2. 反向传播得到放大的梯度
  3. 优化器更新前，把梯度除回来

用法也很固定：

from torch.cuda.amp import GradScaler

scaler = GradScaler()

for data, target in dataloader:
    optimizer.zero_grad()
    
    with autocast():
        output = model(data)
        loss = loss_fn(output, target)
    
    # 用 scaler 包裹 backward
    scaler.scale(loss).backward()
    
    # 用 scaler 包裹 step
    scaler.step(optimizer)
    
    # 更新缩放因子
    scaler.update()

注意最后那个 scaler.update()，很多人会忘。它的作用是动态调整缩放因子——如果连续没出现梯度溢出，就放大一点；如果溢出了，就缩小。说白了就是个自适应策略。

避坑指南： 千万不要手动去调缩放因子！GradScaler 内部有一套成熟的动态调整机制。我曾经手贱去改 init_scale 参数，结果训练直接炸了。让它自己玩就行。

完整的 AMP 训练循环

好，我们把上面两个东西拼起来。这是我在实际项目中用的标准模板：

import torch
from torch.cuda.amp import autocast, GradScaler

model = MyModel().cuda()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
scaler = GradScaler()  # 初始化缩放器

for epoch in range(num_epochs):
    for batch_idx, (data, target) in enumerate(dataloader):
        data, target = data.cuda(), target.cuda()
        optimizer.zero_grad()
        
        # 前向 + loss 用 autocast
        with autocast():
            output = model(data)
            loss = loss_fn(output, target)
        
        # 反向传播用 scaler
        scaler.scale(loss).backward()
        
        # 参数更新用 scaler
        scaler.step(optimizer)
        
        # 更新缩放因子
        scaler.update()
        
        if batch_idx % 100 == 0:
            print(f'Epoch {epoch} Batch {batch_idx} Loss {loss.item():.4f}')

这个循环我用了两年多，基本没出过问题。你可能会问：scaler.step(optimizer) 和普通的 optimizer.step() 有什么区别？其实 scaler.step 内部会先检查梯度有没有溢出，如果有溢出就跳过这步更新，防止模型参数被污染。

AMP 的核心知识体系

为了让你更直观地理解，我画了张图：

AMP 训练循环核心流程 数据加载 (fp32) with autocast(): 前向传播 (自动选择 fp16/fp32) Loss 计算 (自动选择 fp16/fp32) GradScaler 处理 scaler.scale(loss).backward() scaler.step(optimizer) scaler.update() 参数更新完成 autocast: 自动精度选择 GradScaler: 防梯度下溢 update: 动态调整缩放因子

这张图把 AMP 的三个关键步骤串起来了。你照着这个流程写代码，基本不会出错。

实际项目中的注意事项

最后分享几个我踩过的坑：

  • 不要混用不同缩放器：如果你有多个 loss，每个 loss 都要用同一个 scaler 来 scale。我曾经试过每个 loss 单独建一个 scaler，结果梯度更新乱套了。
  • 验证集不需要 autocast：推理时用 torch.no_grad() 就行，不需要 autocast。加了反而可能影响精度。
  • 检查梯度溢出：如果 loss 突然变成 NaN，可以用 scaler.check_overflow() 来诊断。我一般会在训练初期加个判断，一旦溢出就打印警告。
核心总结： AMP 不是黑魔法，它只是把 fp16 和 fp32 的切换自动化了。你只需要记住三句话——前向用 autocast，反向用 scaler，最后别忘了 update。

好了，这一章的内容就这些。代码模板可以直接复制到你的项目里用，我保证它跑得稳。

← 上一章 📖 返回目录 下一章 →