Python中PyTorch如何实现梯度累积_在小批量上叠加梯度实现大Batch

梯度累积为什么不能直接改 batch_size

显存不够用，想上大 batch 又怕 OOM，很多人第一反应是去调 DataLoader 的 batch_size——这条路走不通。PyTorch 的 optimizer.step() 每调一次，梯度就被清空，参数也更新一次。你要是强行塞更大的 batch，要么显存爆掉，要么新梯度把旧梯度覆盖掉，根本谈不上“累积”。

真正有效的做法是：分多次小 batch 做 backward，但只在第 N 次时调用 optimizer.step() 和 optimizer.zero_grad()。这里有几个关键点：

• 每次 loss.backward() 默认会把梯度累加到 .grad 上——前提是没清理过梯度。
• optimizer.zero_grad() 只能在累积开始前或 step 完成后调用，中间绝对不能清。
• loss 必须除以累积步数，否则梯度量级会随着步数线性膨胀，后果很严重。

torch.no_grad() 不参与梯度累积，别误用

代码里最容易踩的坑，是在累积循环里混进 torch.no_grad()。比如有人为了“省显存”，对 validation 做了 no_grad，结果顺手套到了 train loop 里——这样一来，所有 backward() 全部失效，model.parameters()[0].grad 始终是 None。

梯度累积必须全程在梯度计算模式下进行：

• 确保 model.train() 已经调用。
• 不要用 with torch.no_grad(): 包裹任何训练步骤。
• 只有评估、推理、EMA 更新这些非训练逻辑才用 no_grad。

典型错误现象是报 RuntimeError: element 0 of tensors does not require grad and does not ha ve a grad_fn——基本就是某处意外禁用了梯度，排查的时候优先看 no_grad 的位置。

手动控制 optimizer.step() 和 zero_grad() 的时机

核心就三件事：累 loss、累梯度、按周期更新。不需要额外库，也不需要什么钩子，纯靠 if 判断步数就能搞定。

假设目标等效 batch 是 64，当前 GPU 能跑 8：

accumulation_steps = 64 // 8  # = 8
for i, (x, y) in enumerate(dataloader):
    pred = model(x)
    loss = criterion(pred, y) / accumulation_steps  # 关键：缩放 loss
    loss.backward()  # 梯度自动累加
if (i + 1) % accumulation_steps == 0:
    optimizer.step()      # 更新参数
    optimizer.zero_grad() # 清空梯度，准备下一轮累积
loss / accumulation_steps 必须做，否则梯度相当于放大了 8 倍。
zero_grad() 放在 step() 后面，而不是前面；否则刚算完就清空，等于白干。
注意 i + 1：step 是在 batch 结束后触发，不是开头。
如果 dataloader 长度不能被整除，最后一组不足 accumulation_steps 的 batch 也要 step，否则会漏更新。
混合精度训练（torch.cuda.amp）下梯度累积要多一步
用 autocast + GradScaler 时，scaler.scale(loss).backward() 会把梯度乘上 scale 值再累加。好在 scaler.step(optimizer) 内部会自动处理 unscale，但前提是梯度确实被累积了。
容易漏的关键点：
必须用 scaler.unscale_(optimizer) 手动 unscale 梯度，再做梯度裁剪（如 torch.nn.utils.clip_grad_norm_），否则裁剪的是放大后的梯度。
scaler.step() 只应在累积完成时调用，且需传入 optimizer。
别忘了 scaler.update()，否则 scale 值不会自适应调整。
常见错误现象：NaN loss 或训练不稳定，大概率是没 unscale 就裁剪，或者 scaler.update() 被跳过了。
梯度累积本身不改变模型行为，但它把“更新频率”和“数据吞吐节奏”解耦了。实际调试时，最容易被忽略的是 loss 缩放和 zero_grad 的位置——这两处一出错，整个累积就失效，而且很难一眼看出来。