梯度累积(Gradient Accumulation)是一种训练技巧,它可以让我们在显存不那么大的 GPU 上用更大的 batch size 进行训练。
它的思路很自然:以前是在每一个 batch 上都进行参数更新;现在累积几个 batch 的梯度后,再更新参数,变相地增大了 batch size.
介绍 autograd 的那篇文章——Pytorch 中的自动求导与(动态)计算图——提到,Pytorch 的自动求导模块默认将计算好的梯度累加到各个参数的 .grad
属性上。因此在 Pytorch 中进行梯度累积很容易:
optimizer.zero_grad()
accumulation_steps = 10
for i, batch in enumerate(batches):
# Scale the loss
loss = calculate_loss(batch) / accumulation_steps
loss.backward()
if (i + 1) % accumulation_steps == 0:
optimizer.step()
# Reset gradients, for the next accumulated batches
optimizer.zero_grad()
注意这一行:
loss = calculate_loss(batch) / accumulation_steps
为什么要把每一个 batch 的损失函数都除以accumulation_steps
呢?
——一般情况下,损失函数会除以 batch size 求平均。calculate_loss(batch)
求出了当下 batch size 下的平均,为了模拟更大的 batch size,还要除以 accumulation_steps
。因为要模拟的 batch size 大小等于 batch size * accumulation_steps
!
为了方便理解,我画了一张图。考虑两个样本的简单情况。正常情况下,batch size = 2,求出损失函数后要除以 2,然后反向传播求梯度。而在梯度累积的情况下,两个样本分别前向传播得到一个损失函数,并通过反向传播得到了梯度(accumulation steps = 2)。我们要把两个梯度加起来,用于更新参数。为了与正常情况保持一致,单样本的损失函数必须除以 2 (accumulation steps,而非 batch size)
假设在某个模型下,GPU 显存允许的最大 batch size 为 20。通过梯度累积的方式,我们可以实现大小为 20 * accumulation_steps
的 batch size~
先别急着走,现在考虑一种极端情况:在某个大模型下,如果 GPU 连一个样本都跑不了怎么办?——梯度检查 gradient-checkpointing!Pytorch 中的 checkpoint