4. 梯度累积：小显存跑大Batch的秘诀

做深度学习训练，最头疼的事是什么？

对我来说，就是显存不够用。

明明想跑个1024的batch size，结果显卡直接OOM。换小batch吧，模型收敛又慢，效果还差。怎么办？

梯度累积就是来解决这个问题的。

4.1 梯度累积原理

先说说梯度下降的本质。

你想想看，每次前向传播算出loss，然后反向传播算出梯度，最后用梯度更新参数。这是标准流程。

梯度累积干了什么呢？

它把「算梯度」和「更新参数」这两步拆开了。

具体来说：

• 先跑N个小batch，每个batch都算梯度
• 但这些梯度不立即用来更新参数
• 而是累加起来
• 等累够了N步，再用累积的梯度去更新一次参数

说白了，就是用时间换空间。

你本来一次想跑1024的batch，显存放不下。现在拆成4个256的小batch，每个都能跑。跑完4次，等效于跑了一次1024。

我在项目中遇到过很多次这种情况。有一次做BERT预训练，目标batch size是8192，单卡根本跑不了。用梯度累积，拆成16步，每步512，完美解决。

4.2 梯度累积实现

实现起来其实不复杂。我直接给代码，你们一看就懂。

# PyTorch实现梯度累积
accumulation_steps = 4  # 累积步数
optimizer.zero_grad()   # 注意：只清零一次

for i, (inputs, labels) in enumerate(dataloader):
    outputs = model(inputs)
    loss = criterion(outputs, labels)
    
    # 关键：除以累积步数，保持loss尺度一致
    loss = loss / accumulation_steps
    loss.backward()
    
    # 每accumulation_steps步更新一次
    if (i + 1) % accumulation_steps == 0:
        optimizer.step()
        optimizer.zero_grad()

嗯，这里要注意几个细节：

• loss要除以累积步数：不然累积后的梯度会变大，相当于学习率被放大了
• optimizer.zero_grad()的位置：只在更新参数后清零，不要每步都清
• 最后一个batch的处理：如果数据量不是累积步数的整数倍，记得处理剩余部分

4.3 梯度累积与Batch Size关系

很多人问我：梯度累积真的等价于增大batch size吗？

答案是：不完全等价，但非常接近。

为什么说不完全等价？

因为BN层（Batch Normalization）的存在。

BN层在计算均值和方差时，用的是当前batch的统计量。如果你用256的batch size累积4步，BN看到的还是256的数据，而不是1024。

那怎么办？

我有两个建议：

1. 使用SyncBN：跨卡同步BN统计量，但单卡场景下用不了
2. 增大单次batch size：在显存允许范围内，尽量把单次batch设大一点

场景	单次Batch	累积步数	有效Batch	BN效果
小显存	64	8	512	较差
中等显存	128	4	512	中等
大显存	256	2	512	较好

从表格能看出来：单次batch越大，BN效果越好。所以我的建议是——优先增大单次batch，不够了再用累积补。

4.4 梯度累积实战案例

讲个真实案例吧。

去年我做了一个图像分类项目，用的是EfficientNet-B4。目标batch size是512，但我的RTX 3090只能跑128。

我的配置是这样的：

# 实战配置
config = {
    'batch_size': 128,          # 单次batch
    'accumulation_steps': 4,    # 累积4步
    'effective_batch': 512,     # 等效batch
    'learning_rate': 0.001,     # 按512batch调优的学习率
    'warmup_steps': 1000        # 预热步数
}

训练过程中，我特别注意了三点：

• 学习率调整：有效batch变大了，学习率也要相应调大。我一般按线性缩放规则来
• 梯度裁剪：累积梯度可能会让梯度爆炸，加上梯度裁剪更安全
• 日志记录：每步都记录loss，但只在更新参数时打印，不然日志太密集

最后的结果：

• 显存占用：从OOM降到8.2GB
• 训练时间：增加了约15%（因为多了一些调度开销）
• 模型精度：和直接跑512batch几乎一样，top-1准确率差不到0.1%

这个案例说明什么？

梯度累积不是银弹，但它确实是小显存用户的救星。只要配置得当，效果和直接跑大batch几乎没区别。

这张图把流程说得很清楚了。小batch跑N次，梯度累加，最后一次性更新。显存没变，但等效batch变大了。