1. 为什么我们需要梯度累积当你在训练深度学习模型时可能会遇到一个令人头疼的问题显存不够用。特别是当模型越来越大或者你想尝试更大的batch size时显存限制就成了拦路虎。这时候梯度累积Gradient Accumulation就像是一个救星它能让你在有限的显存下变相扩大batch size。我刚开始用PyTorch训练模型时就经常被显存不足的问题困扰。比如我想用batch size128训练一个ResNet模型但我的GPU只能承受batch size32。这时候梯度累积就派上用场了。它的核心思想很简单把多个小batch的梯度累积起来等累积到足够数量后再一次性更新模型参数。举个例子假设你希望等效的batch size是128但实际显存只能支持batch size32。那么你可以用batch size32训练4个batch把这4个batch的梯度累积起来最后用累积的梯度更新一次参数这样虽然每次前向传播和反向传播的batch size还是32但参数更新的效果相当于batch size128。我在实际项目中多次使用这个技巧效果确实不错特别是当显存紧张但又想保持较大batch size时。2. 梯度累积的工作原理2.1 传统训练 vs 梯度累积训练传统的训练方式是每个batch都更新一次参数for data, target in train_loader: optimizer.zero_grad() output model(data) loss criterion(output, target) loss.backward() optimizer.step() # 每个batch都更新参数而梯度累积的训练方式是这样的accumulation_steps 4 # 累积4个batch的梯度 for i, (data, target) in enumerate(train_loader): output model(data) loss criterion(output, target) loss loss / accumulation_steps # 损失标准化 loss.backward() if (i 1) % accumulation_steps 0: optimizer.step() # 累积够4个batch才更新参数 optimizer.zero_grad()关键区别在于不是每个batch都调用optimizer.step()需要把loss除以累积步数因为PyTorch的backward()是梯度累加而不是平均只在累积够指定步数后才更新参数和清零梯度2.2 梯度累积的数学原理从数学上看梯度累积相当于对多个batch的梯度求平均。假设我们要累积k个batch每个batch计算出的梯度是∇Lᵢ累积后的总梯度是(∇L₁ ∇L₂ ... ∇Lₖ)/k用这个平均梯度来更新参数这就是为什么我们要把loss除以accumulation_steps - 这样最终的梯度就是多个batch梯度的平均值而不是简单的累加。3. PyTorch中的梯度累积实现3.1 基础实现代码下面是一个完整的PyTorch梯度累积实现示例model MyModel().to(device) optimizer torch.optim.Adam(model.parameters(), lr0.001) criterion nn.CrossEntropyLoss() accumulation_steps 4 # 累积4个batch batch_size 32 # 实际batch size effective_batch_size batch_size * accumulation_steps # 等效batch size128 for epoch in range(num_epochs): model.train() for i, (inputs, labels) in enumerate(train_loader): inputs inputs.to(device) labels labels.to(device) # 前向传播 outputs model(inputs) loss criterion(outputs, labels) # 标准化损失并反向传播 loss loss / accumulation_steps loss.backward() # 累积够步数后更新参数 if (i 1) % accumulation_steps 0: optimizer.step() optimizer.zero_grad() # 可以在这里添加验证或其他操作 if (i 1) % evaluation_steps 0: evaluate_model()3.2 实现中的注意事项学习率调整因为等效batch size变大了通常需要相应增大学习率。我一般会按累积步数的平方根比例调整比如累积4个batch学习率可以乘以2。BatchNorm层如果你模型中有BatchNorm层要注意它看到的是实际的batch size而不是等效的batch size。这种情况下你可能需要调整BatchNorm的momentum参数。梯度裁剪使用梯度累积时梯度可能会变得比较大建议添加梯度裁剪torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm1.0)混合精度训练梯度累积可以和混合精度训练很好地结合使用进一步节省显存scaler torch.cuda.amp.GradScaler() with torch.cuda.amp.autocast(): outputs model(inputs) loss criterion(outputs, labels) loss loss / accumulation_steps scaler.scale(loss).backward() if (i 1) % accumulation_steps 0: scaler.step(optimizer) scaler.update() optimizer.zero_grad()4. 梯度累积的性能分析与调优4.1 性能对比实验我在实际项目中做过对比实验使用ResNet-18在CIFAR-10数据集上训练方式Batch Size显存占用训练时间/epoch最终准确率普通训练12810.2GB45s92.5%普通训练323.1GB50s91.3%梯度累积32(等效128)3.1GB55s92.1%可以看到梯度累积在几乎不增加显存占用的情况下达到了接近大batch size训练的效果虽然训练时间稍长一些。4.2 调优建议累积步数选择不是越大越好。我一般建议累积2-8步太多会导致参数更新太不频繁可能影响收敛。学习率调整可以尝试线性缩放规则(linear scaling rule) - 如果batch size扩大k倍学习率也扩大k倍。或者更保守的平方根缩放(sqrt scaling) - 学习率扩大√k倍。warmup策略使用大batch size(即使是等效的)时配合学习率warmup通常效果更好def adjust_learning_rate(optimizer, epoch, warmup_epochs5): if epoch warmup_epochs: lr base_lr * (epoch 1) / warmup_epochs else: lr base_lr for param_group in optimizer.param_groups: param_group[lr] lr验证频率因为参数更新变少了可以适当增加验证频率比如每累积更新2-3次就验证一次。不同层的累积对于特别大的模型可以尝试对不同部分使用不同的累积策略。比如视觉部分的梯度累积4次文本部分累积2次。在实际项目中我发现梯度累积特别适合以下场景模型很大显存紧张想要使用大的batch size但硬件不支持做对比实验时需要保持batch size一致在预训练大模型时配合混合精度使用