PyTorch中的梯度累积

我们在训练神经网络的时候，超参数batch_size的大小会对模型最终效果产生很大的影响，通常的经验是，batch_size越小效果越差；batch_size越大模型越稳定。理想很丰满，现实很骨感，很多时候不是你想增大batch_size就能增大的，受限于显存大小等因素，我们的batch_size往往只能设置为2或4，否则就会出现"CUDA OUT OF MEMORY"(OOM)报错。如何在有限的计算资源下，采用更大的batch_size进行训练，或者达到和大batch_size一样的效果？这就是梯度累加（Gradient Accumulation）技术了

以PyTorch为例，正常来说，一个神经网络的训练过程如下：

for idx, (x, y) in enumerate(train_loader):
    pred = model(x)
    loss = criterion(pred, y)
    
    loss.backward()
    optimizer.step()
    optimizer.zero_grad()
    
    if (idx+1) % eval_steps == 0:
        eval()

如果你想设置batch_size=64结果爆显存了，那么不妨设置batch_size=16，然后定义一个变量accum_steps=4，每个mini-batch仍然正常前向传播以及反向传播，但是反向传播之后并不进行梯度清零，因为PyTorch中的loss.backward()执行的是梯度累加的操作，所以当你调用4次loss.backward()后，这4个mini-batch的梯度都会累加起来。但是，我们需要的是一个平均的梯度，或者说平均的损失，所以我们应该将每次计算得到的loss除以accum_steps

accum_steps = 4

for idx, (x, y) in enumerate(train_loader):
    pred = model(x)
    loss = criterion(pred, y)
    
    # normlize loss to account for batch accumulation
    loss = loss / accum_steps
    
    loss.backward()
    
    if (idx+1) % accum_steps == 0 or (idx+1) == len(train_loader):
        optimizer.step()
        optimizer.zero_grad()
        if (idx+1) % eval_steps:
            eval()

总的来说，梯度累加就是计算完每个mini-batch的梯度后不清零，而是做梯度的累加，当累加到一定的次数之后再更新网络参数，然后将梯度清零。通过这种延迟更新的手段，可以实现与采用大batch_size相近的效果

References

• pytorch中的梯度累加(Gradient Accumulation)
• Gradient Accumulation in PyTorch
• PyTorch中在反向传播前为什么要手动将梯度清零？