查看全课大纲（第 38 / 40 节）

PyTorch 2.x 训练循环与混合精度

小象实战讲义 · 深度学习基础：从神经网络到 Transformer

模型结构讲完，最后一步是把它们真正训练出来。本节讲 PyTorch 2.x 的训练循环主干——自动求导、六步训练循环、自动混合精度（AMP）与梯度累积，重点是 AMP 为什么既能省显存又不掉精度，以及 2.x 里该用哪套 API。本节是第 9 章《工程实践与收口》的开篇。先修：前向/反向传播与梯度下降（第 2 章）。后续：下一节讲实验管理与复现，再下一节讲推理部署与量化蒸馏。

💡 核心导读

自动求导（Autograd）：PyTorch的核心，通过动态计算图自动计算梯度，是训练得以进行的基础。
标准训练循环：包含数据加载、模型前向、损失计算、梯度清零、反向传播、参数更新六个关键步骤。
混合精度训练（AMP）：通过让部分计算（如前向传播）在float16（半精度）下进行，减少显存占用和加速计算，同时使用float32（单精度）维护主权重以保证数值稳定性。
梯度累积：一种在有限显存下模拟更大批次训练的实用技巧，通过多次小批次的前后向传播累积梯度，再一次性更新参数。
PyTorch 2.x 新特性：引入了更简洁、高效的API（如torch.compile），但本节聚焦于稳定且核心的训练循环模式。

🧭 前置知识

矩阵乘法与张量形状的基本规则。
损失函数（如交叉熵）和梯度下降法的基本概念。
链式法则：理解反向传播中梯度是如何从输出层传递回输入层的。

动态计算图与自动求导（Autograd）

PyTorch与早期一些框架（如静态图的TensorFlow 1.x）的核心区别在于其动态计算图。计算图在代码运行时动态构建，每次前向传播都可能生成一个不同的图结构，这使得编写包含条件分支、循环的模型变得非常直观。

自动求导系统（torch.autograd）是这一切的引擎。当你对张量设置 requires_grad=True 时，PyTorch会开始跟踪在其上执行的所有操作，形成一个由函数（Function）节点和张量（Tensor）节点组成的有向无环图（DAG）。前向传播执行操作并记录梯度计算所需的信息。反向传播时，只需对最终的标量损失调用 .backward()，autograd 就会依据链式法则，自动计算图中所有 requires_grad=True 的张量的梯度，并累积到它们的 .grad 属性中。

import torch

torch.manual_seed(0)

# 1. 创建需要梯度的张量
x = torch.tensor([2.0, 3.0], requires_grad=True)
w = torch.tensor([1.0, 2.0], requires_grad=True)
b = torch.tensor(0.5, requires_grad=True)

# 2. 前向传播（动态构建计算图）
y = w * x + b  # y = [1*2, 2*3] + 0.5 = [2.5, 6.5]
z = y.sum()    # z = 2.5 + 6.5 = 9.0

# 3. 反向传播，自动计算梯度
z.backward()

# 4. 查看梯度
print("梯度 dz/dx =", x.grad.tolist())  # dz/dx = w = [1.0, 2.0]
print("梯度 dz/dw =", w.grad.tolist())  # dz/dw = x = [2.0, 3.0]
print("梯度 dz/db =", b.grad.item())    # dz/db = 2（y 有两个元素，各贡献 1）

梯度 dz/dx = [1.0, 2.0]
梯度 dz/dw = [2.0, 3.0]
梯度 dz/db = 2.0

关键点：z.backward() 计算的是标量 z 对图中所有 requires_grad=True 的叶节点（即用户直接创建的张量，如 x, w, b）的梯度。梯度会累积在 .grad 属性中，因此在每次参数更新前需要手动清零。

PyTorch 标准训练循环剖析

一个完整的训练循环包含数据准备、模型定义、损失函数、优化器以及迭代训练几个部分。老版本 PyTorch 里还会遇到 Variable 包装，它早已并入 Tensor，现在直接写 Tensor 就行。

下面我们用一个极简的两层全连接网络在合成数据上演示标准流程。

import torch
import torch.nn as nn
import torch.optim as optim

torch.manual_seed(0)

# --- 1. 合成数据准备 (Batch Size=4, 特征维度=10, 3个类别) ---
batch_size, d_in, d_hidden, d_out = 4, 10, 32, 3
x = torch.randn(batch_size, d_in)
y = torch.randint(0, d_out, (batch_size,))

# --- 2. 定义模型 ---
class SimpleMLP(nn.Module):
    def __init__(self, input_dim, hidden_dim, output_dim):
        super().__init__()
        self.layer1 = nn.Linear(input_dim, hidden_dim)
        self.relu = nn.ReLU()
        self.layer2 = nn.Linear(hidden_dim, output_dim)

    def forward(self, x):
        x = self.layer1(x)
        x = self.relu(x)
        x = self.layer2(x)
        return x

model = SimpleMLP(d_in, d_hidden, d_out)

# --- 3. 定义损失函数和优化器 ---
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.01)

# --- 4. 单个训练步骤（Step） ---
def train_one_step(model, x, y, criterion, optimizer):
    # 4.1 前向传播
    logits = model(x)  # 形状: (4, 3)
    loss = criterion(logits, y)

    # 4.2 梯度清零（重要！防止梯度累积）
    optimizer.zero_grad()

    # 4.3 反向传播
    loss.backward()

    # 4.4 参数更新
    optimizer.step()

    return loss.item()

# 执行两步训练，观察损失变化
print("训练前模型第一层权重前两个值:", [round(v, 4) for v in model.layer1.weight.data[0, :2].tolist()])
loss1 = train_one_step(model, x, y, criterion, optimizer)
print(f"第一步训练后损失: {loss1:.4f}")
print("第一步后模型第一层权重前两个值:", [round(v, 4) for v in model.layer1.weight.data[0, :2].tolist()])

# 用同样的数据再训练一步（仅作演示，实际应用会用新数据）
loss2 = train_one_step(model, x, y, criterion, optimizer)
print(f"第二步训练后损失: {loss2:.4f}")

训练前模型第一层权重前两个值: [0.1431, 0.1272]
第一步训练后损失: 1.0386
第一步后模型第一层权重前两个值: [0.1427, 0.1272]
第二步训练后损失: 1.0229

注意：由于学习率很小且只训练两步，权重变化可能微乎其微，但损失计算和梯度流程已完整走通。

训练循环六步法：

前向传播：model(x) 计算预测值。
计算损失：criterion(pred, y) 计算标量损失。
梯度清零：optimizer.zero_grad() 将模型参数 .grad 属性置零。
反向传播：loss.backward() 自动计算各参数梯度。
参数更新：optimizer.step() 根据梯度（如SGD规则）更新参数。
循环：对数据加载器 (DataLoader) 中的每个批次重复步骤1-5。

自动混合精度训练（AMP）

混合精度训练是加速训练、节省显存的关键技术。其核心思想是：使用 float16（半精度）进行张量计算和存储，因为 float16 所需显存仅为 float32 的一半，且在支持Tensor Core的现代GPU上计算速度更快；同时，保留一份 float32 的“主权重”副本用于参数更新，以保持训练的数值稳定性。

PyTorch 2.x 把 AMP 统一到了 torch.amp：torch.amp.autocast(device_type, dtype) 负责前向时逐个算子挑精度，torch.amp.GradScaler(device_type) 负责梯度缩放。老的 torch.cuda.amp.autocast() / torch.cuda.amp.GradScaler() 从 2.4 起进入弃用路径，新代码应显式传 device_type。CPU 上没有 float16 的 Tensor Core，但 bfloat16 的 autocast 可用，下面的演示就在 CPU 上真跑。

import torch
import torch.nn as nn
import torch.optim as optim
import torch.amp as amp

torch.manual_seed(0)

batch_size, d_in, d_hidden, d_out = 4, 10, 32, 3
x = torch.randn(batch_size, d_in)
y = torch.randint(0, d_out, (batch_size,))

model = nn.Sequential(nn.Linear(d_in, d_hidden), nn.ReLU(), nn.Linear(d_hidden, d_out))
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=1e-3)

# 1) 全 float32 前向
out_fp32 = model(x)

# 2) autocast 下前向：算子精度由 PyTorch 自动挑选（CPU 支持 bfloat16）
with torch.autocast(device_type="cpu", dtype=torch.bfloat16):
    out_bf16 = model(x)
print("autocast 关闭时输出 dtype:", out_fp32.dtype)
print("autocast 开启时输出 dtype:", out_bf16.dtype)
print(f"两者最大绝对差: {(out_fp32 - out_bf16.float()).abs().max().item():.6f}")

# 3) GradScaler：把损失放大后再反向，更新前自动反缩放
scaler = amp.GradScaler("cpu")
optimizer.zero_grad()
with torch.autocast(device_type="cpu", dtype=torch.bfloat16):
    loss = criterion(model(x), y)
scaler.scale(loss).backward()
print(f"原始损失: {loss.item():.4f}")
print(f"缩放后损失: {scaler.scale(loss).item():.4f}")
scaler.step(optimizer)
scaler.update()
print(f"当前缩放因子: {scaler.get_scale():.1f}")

autocast 关闭时输出 dtype: torch.float32
autocast 开启时输出 dtype: torch.bfloat16
两者最大绝对差: 0.001581
原始损失: 1.0386
缩放后损失: 68068.1328
当前缩放因子: 65536.0

为什么需要 GradScaler？ float16 的表示范围（约 5.96e-8 ~ 65504）远小于 float32。在深度网络中，梯度值可能非常小，在 float16 中会下溢为零，导致训练停滞。GradScaler 将损失乘以一个缩放因子（如 2**16），使得梯度在 float16 的表示范围内，然后在优化器更新前再将梯度除以相同的因子。

能省多少显存？ 权重和激活用 fp16 存，理论减半；但 Adam 的动量与方差仍要 fp32（每参数两份），所以整体峰值显存通常只降 20%–40%。真正吃显存的是激活值，因此大 batch 训练里 AMP 的收益最明显——这也是它常和梯度累积一起用的原因。

float16 还是 bfloat16？

格式	指数位	尾数位	动态范围	是否需要 GradScaler
float32	8	23	约 ±3.4e38	——
float16	5	10	约 6e-8 ~ 65504	需要
bfloat16	8	7	与 float32 同量级	通常不需要

float16 的毛病就在那个 5 位指数：范围太窄，梯度稍小就下溢成 0、稍大就溢出成 inf。bfloat16 砍掉尾数精度、保留与 float32 相同的 8 位指数，因此几乎不溢出，也就不必配 GradScaler；代价是尾数只剩 7 位（大约 2–3 位十进制有效数字）。实践中：Ampere 及以后的 GPU 两种都支持，训练大模型常优先 bf16；V100/T4 这类老卡只有 fp16，必须配 GradScaler。

autocast 到底把什么降了精度？

autocast 不是“把整张计算图都变成 fp16”，而是逐算子决策：

降到低精度：矩阵乘、卷积、线性层——它们是算力大头，且对精度不敏感。
保持 float32：softmax、LayerNorm、损失函数、各类归约——它们对数值范围敏感，降精度容易出 NaN。

这就是“混合”二字的来历：同一张图里两种精度共存，PyTorch 负责在算子边界插入 cast。所以常见错误表里强调「不要手动 model.half()」——那会把该保持 fp32 的部分也一起压掉。

GradScaler 的三个动作

scale：反向传播前把损失乘以缩放因子 $S$ （初值 $2^{16}$ ），梯度整体放大，避开下溢。
unscale + step：更新前把梯度除回 $S$ ；若发现 inf/nan，就跳过这一步、不污染权重。
update：连续若干步没溢出就把 $S$ 翻倍，一旦溢出就减半，动态贴着当前数值范围走。

和 torch.compile 的关系

torch.compile 负责把训练循环编译成融合 kernel、削掉 Python 与 kernel 启动开销；AMP 负责降低单次计算的数据位宽。两者正交、常一起用：外层 compile、内层 autocast。compile 自己有坑（动态形状导致重编译、图断裂回退 eager），本节不展开，但要知道它和 AMP 解决的是两个不同层面的问题。

梯度累积：突破显存限制的技巧

当可用的GPU显存无法容纳理想大小的批次（batch）时，梯度累积（Gradient Accumulation）是一种有效的解决方案。其原理是：在多个小批次（micro-batch）上依次进行前向传播和反向传播，但不立即更新参数，而是让梯度累积在参数的 .grad 属性中。在累积了 N 个小批次后，再执行一次 optimizer.step() 和 optimizer.zero_grad()。这等效于用 N * micro_batch_size 的大批次进行了一次训练，但峰值显存占用仅为一个小批次的大小。

import torch
import torch.nn as nn
import torch.optim as optim

torch.manual_seed(0)

# 模拟一个简单模型
model = nn.Linear(10, 3)
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.01)

# 模拟数据：假设我们想用 batch_size=8 训练，但显存只够 micro_batch_size=2
micro_batch_size = 2
accumulation_steps = 4  # 累积4步，等效 batch_size=8
total_samples = 8

# 生成8个样本
x_all = torch.randn(total_samples, 10)
y_all = torch.randint(0, 3, (total_samples,))

# 模拟梯度累积训练一步（等效于用8个样本训练一次）
optimizer.zero_grad()  # 在累积开始前清零梯度

total_loss = 0.0
for i in range(accumulation_steps):
    # 取一个小批次
    start = i * micro_batch_size
    end = start + micro_batch_size
    x_micro = x_all[start:end]
    y_micro = y_all[start:end]

    # 前向传播
    logits = model(x_micro)
    loss = criterion(logits, y_micro)
    total_loss += loss.item()

    # 反向传播（梯度会累积到 model.parameters() 的 .grad 中）
    # 注意：这里损失要除以累积步数，使得累积梯度的量级与大批次训练时一致
    (loss / accumulation_steps).backward()

# 累积完成后，更新参数
optimizer.step()
# 注意：下一轮累积开始前需要再次调用 optimizer.zero_grad()

print(f"梯度累积完成。模拟的等效大批次损失约为: {total_loss / accumulation_steps:.4f}")
print(f"参数已更新一次。")

梯度累积完成。模拟的等效大批次损失约为: 1.3827
参数已更新一次。

梯度累积的三个坑

BatchNorm 的统计量按 micro-batch 更新：累积 N 步虽然等效大批次，但 BN 的 running mean/var 是每个小批次各更新一次，统计量与真正的大批次不一致。要么改用 LayerNorm/GroupNorm，要么用同步 BN。
学习率调度器要跟着 optimizer.step() 走：scheduler 一般按 step 计步，累积时一次 step 覆盖 N 个 micro-batch，scheduler.step() 必须与 optimizer.step() 放在同一位置，否则学习率衰减会快 N 倍。
Dropout 掩码不共享：大批次是一次采样，累积是 N 次独立采样，正则效果略有差异——通常影响不大，但小数据上会看到 loss 曲线更抖。

这三条都属于「等效大批次」名不副实的地方，调参时要心里有数。

📝 动手练一练

问题：在标准训练循环中，如果忘记调用 optimizer.zero_grad() 会导致什么问题？请编写一小段代码验证你的想法。

👉 点击查看参考答案

思路：optimizer.zero_grad() 的作用是清空模型参数的梯度缓存。如果忘记调用，每次 loss.backward() 计算出的梯度会与上一轮残留的梯度累加，导致参数更新方向错误，训练无法收敛。

import torch
import torch.nn as nn
import torch.optim as optim

torch.manual_seed(0)

# 一个极简的线性模型 y = w*x
model = nn.Linear(1, 1, bias=False)  # 只有一个参数 w
optimizer = optim.SGD(model.parameters(), lr=0.1)
criterion = nn.MSELoss()

# 固定数据
x = torch.tensor([[1.0]])
y = torch.tensor([[2.0]])

print("初始权重 w =", round(model.weight.data.item(), 4))

# 第一次训练（正常）
optimizer.zero_grad()
loss1 = criterion(model(x), y)
loss1.backward()
print("第一次反向传播后梯度 =", round(model.weight.grad.item(), 4))
optimizer.step()
print("第一次更新后权重 w =", round(model.weight.data.item(), 4))

# 第二次训练（忘记 zero_grad）
# optimizer.zero_grad()  # 故意注释掉
loss2 = criterion(model(x), y)
loss2.backward()
print("第二次反向传播后梯度（累积后） =", round(model.weight.grad.item(), 4))
optimizer.step()
print("第二次更新后权重 w =", round(model.weight.data.item(), 4))

初始权重 w = -0.0075
第一次反向传播后梯度 = -4.015
第一次更新后权重 w = 0.394
第二次反向传播后梯度（累积后） = -7.227
第二次更新后权重 w = 1.1167

分析：第一次梯度是 -4.015（= 2(w−2)，w 初始 -0.0075），更新后 w 变成 0.394。第二次若正常清零，梯度应只有 2(0.394−2) = -3.212；但没清零，实测累积成 -7.227（≈ 两次之和），优化器按 -7.227 迈步，w 从 0.394 一步跳到 1.117——已经越过目标值 2 的一半，再多几步就发散。这就是「梯度累积」在训练循环里必须显式清零的原因。

常见错误

错误理解或写法	现象或报错	原因	正确做法
忘记在每轮迭代前调用 `optimizer.zero_grad()`。	训练损失剧烈震荡或发散，不收敛。	梯度在 `.grad` 属性中累积，导致每次参数更新使用了多轮梯度的和，步长和方向错误。	在每个训练步骤的开始或 `loss.backward()` 之前，调用 `optimizer.zero_grad()`。
在 `validation` 或 `test` 阶段未使用 `torch.no_grad()`。	显存占用异常增长，甚至爆显存；计算速度变慢。	PyTorch 默认会为 `requires_grad=True` 的张量构建计算图以支持求导。在推理时这是不必要的，会浪费资源。	使用 `with torch.no_grad():` 上下文管理器包裹推理代码。
在混合精度训练中，手动将模型或数据转换为 `float16`。	可能引发数值不稳定（NaN），或无法享受AMP的自动类型转换优化。	`autocast()` 上下文管理器会智能地为每个算子选择最佳精度。手动转换会破坏这一机制。	信任 `autocast()`，让它在适合的层（如卷积、线性层）使用 `float16`，在需要高精度的层（如softmax）自动使用 `float32`。
梯度累积时，没有对每个小批次的损失做缩放（`loss/accum_steps`）。	梯度量级被放大 accumulation_steps 倍，等效学习率失控，带动量的优化器与调度器行为异常。	累积 N 步的梯度总和相当于用 N 倍大批次算一次梯度；不缩放就多乘了 N。	写 `(loss / accumulation_steps).backward()`，让梯度量级与大批次一致。

本章小结

PyTorch的动态计算图和自动求导（Autograd） 是模型训练的基础，它自动计算梯度，解放了开发者。
一个标准的训练循环包含：数据加载、前向传播、损失计算、梯度清零、反向传播、参数更新六个核心步骤。
自动混合精度训练（AMP） 通过结合 float16 的计算效率和 float32 的数值稳定性，并借助 GradScaler 管理梯度缩放，能显著提升训练速度并降低显存消耗。
梯度累积是一种在有限显存下模拟大批次训练的有效技巧，通过多次前反向传播累积梯度后再更新参数。
注意训练中的常见陷阱，如忘记梯度清零、在推理阶段未禁用梯度计算等。

📋 行动清单

在本地环境（如有GPU）中，尝试将本节的标准训练循环代码改造成支持AMP的版本，并观察显存占用和训练速度的变化。
编写一个简单的训练脚本，故意省略 optimizer.zero_grad()，观察训练损失曲线的异常情况。
查阅 PyTorch 官方文档中关于 torch.amp 的教程，了解 GradScaler.update() 如何动态调整缩放因子。