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 相近的效果

PyTorch 中的梯度累积

July 27, 2021 • Read: 4953 • Deep Learning • 阅读设置

References