梯度累积实战：小batch size下的训练稳定性提升方案

梯度累积实战：小batch size下的训练稳定性提升方案

在实际工程实践中，受限于显存资源，我们经常需要使用较小的batch size进行训练。然而，小batch size会导致梯度估计不稳定，影响模型收敛和最终性能。

问题复现

以ResNet50为例，使用batch_size=8训练时，损失曲线波动剧烈：

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

# 模拟小batch size训练
model = torchvision.models.resnet50(pretrained=False)
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
criterion = nn.CrossEntropyLoss()

for epoch in range(5):
    for batch_idx, (data, target) in enumerate(train_loader):
        optimizer.zero_grad()
        output = model(data)
        loss = criterion(output, target)
        loss.backward()
        optimizer.step()
        
        if batch_idx % 100 == 0:
            print(f'Epoch: {epoch}, Batch: {batch_idx}, Loss: {loss.item()}')

梯度累积解决方案

通过梯度累积，我们可以在保持小batch size的同时，模拟大batch size的训练效果：

# 梯度累积实现
accumulation_steps = 4
model.train()
for epoch in range(5):
    for batch_idx, (data, target) in enumerate(train_loader):
        output = model(data)
        loss = criterion(output, target) / accumulation_steps  # 归一化
        loss.backward()  # 累积梯度
        
        if (batch_idx + 1) % accumulation_steps == 0:
            optimizer.step()  # 执行优化
            optimizer.zero_grad()
            print(f'Epoch: {epoch}, Step: {batch_idx//accumulation_steps}, Loss: {loss.item()*accumulation_steps}')

性能对比

在CIFAR-10数据集上测试，batch_size=4，累积步数分别为1和4：

结果表明，在保持相同显存占用的前提下，梯度累积可提升模型收敛稳定性，降低最终损失值约15%。