我们在训练神经网络的时候,超参数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
相近的效果