📚 课程目录第4章梯度累积

4. 梯度累积

说实话，刚入行那会儿，我特别羡慕那些有8块A100的团队。他们可以随便把batch size开到1024，而我手里只有一块2080Ti，8G显存，跑个ResNet50 batch size开到32就爆显存了。

后来我学到了梯度累积，这玩意儿简直就是我们穷人的救星。说白了，就是用时间换空间——多跑几步前向传播，攒够了梯度再更新一次参数。

梯度累积的原理

先想想正常的训练流程：

前向传播，算loss
反向传播，算梯度
优化器更新参数

梯度累积改成了这样：

前向传播，算loss
反向传播，算梯度（但不更新参数）
重复步骤1-2，把梯度累加起来
累积够N步后，用累积的梯度更新参数

嗯，这里要注意：累积的是梯度，不是loss。我见过有人直接把loss加起来，然后反向传播一次，这完全不对。loss是标量，梯度是张量，两者性质不同。

核心公式：

假设你想模拟batch size = 1024的效果，但显存只够跑batch size = 128。

那就设累积步数 = 1024 / 128 = 8。

每8步更新一次参数，等效于batch size = 1024。

为什么能等效？

你想想看，正常的batch size = 1024，是一次性算1024个样本的梯度，然后更新。

梯度累积是分8次，每次算128个样本的梯度，加起来，再更新。

数学上，梯度的线性叠加性质保证了这两者是等价的。说白了就是：

grad_total = grad_1 + grad_2 + ... + grad_8

只要每次的样本是独立采样的，结果就没区别。

我的经验：实际训练中，由于Batch Normalization层的存在，两者会有细微差别。BN的均值和方差在累积模式下是分步计算的，不是全局的。所以如果你用BN，建议把累积步数设小一点，或者用SyncBN。

代码实现

我直接给你一个完整的PyTorch实现，这是我项目里一直在用的模板：

import torch
import torch.nn as nn
from torch.utils.data import DataLoader

# 配置
accumulation_steps = 8  # 累积步数
effective_batch_size = 1024  # 目标batch size
actual_batch_size = 128  # 实际能跑的batch size

model = YourModel()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)
scaler = torch.cuda.amp.GradScaler()  # 混合精度配合使用

model.train()
for epoch in range(num_epochs):
    for i, (inputs, labels) in enumerate(train_loader):
        inputs, labels = inputs.cuda(), labels.cuda()
        
        # 混合精度前向
        with torch.cuda.amp.autocast():
            outputs = model(inputs)
            loss = criterion(outputs, labels)
            # 关键：loss要除以累积步数
            loss = loss / accumulation_steps
        
        # 反向传播，累积梯度
        scaler.scale(loss).backward()
        
        # 每accumulation_steps步更新一次
        if (i + 1) % accumulation_steps == 0:
            # 梯度裁剪（可选）
            scaler.unscale_(optimizer)
            torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
            
            # 更新参数
            scaler.step(optimizer)
            scaler.update()
            optimizer.zero_grad()

我曾经踩过的坑：

忘记把loss除以accumulation_steps。如果不除，累积的loss会变成原来的N倍，梯度爆炸。
优化器更新后忘记zero_grad。梯度会一直累积下去，永远不会清零。
最后一个batch不够累积步数时，梯度没更新就结束了。建议用drop_last=True丢弃不完整的batch。

配合混合精度使用

梯度累积和混合精度是绝配。为什么？

混合精度用FP16存梯度，显存直接减半。梯度累积又让你能用小batch模拟大batch。两者一结合，我曾在4G显存的卡上跑出了batch size=512的效果。

具体做法就是上面代码里那样：用GradScaler管理梯度缩放，用autocast做自动混合精度。注意scaler.step()和scaler.update()要放在累积步数判断里面。

实际效果对比

配置	显存占用	等效Batch Size	训练速度
直接跑batch=1024	24GB	1024	1x
batch=128 + 累积8步	4GB	1024	1.1x（略慢）
batch=128 + 累积8步 + 混合精度	2.5GB	1024	0.9x（反而更快）

看到了吗？混合精度不仅省显存，因为FP16计算更快，有时候反而能提速。

什么时候不该用梯度累积？

也不是所有场景都适合。我总结了几种情况：

用Batch Normalization时：累积步数别超过4，否则BN的统计量偏差太大
强化学习：策略梯度方法对梯度延迟敏感，累积步数建议≤2
GAN训练：生成器和判别器交替更新，累积步数要一致，否则训练不稳定

我的建议：刚开始用梯度累积时，先设accumulation_steps=2或4试试。看看loss曲线和正常训练是否一致。没问题了再逐步加大。别一上来就设32，出了问题很难排查。

用SVG看梯度累积流程

下面这张图帮你理清整个流程：

这张图里，红色箭头是循环路径，绿色箭头是更新路径。每次反向传播后，梯度都累加到同一个缓冲区里。直到累积够N步，才触发优化器更新。

嗯，这里有个细节：清零梯度一定要在更新之后。我见过有人把顺序搞反了，结果梯度永远都是0，模型根本不学习。

好了，梯度累积就讲这么多。说白了就是一句话：显存不够，步数来凑。配合混合精度，你甚至能在笔记本上跑出实验室级别的batch size。