关于梯度累积的说明

Question

我正在尝试更好地了解梯度累积的工作原理以及它为何有用。为此，我想问一下这两种可能的具有梯度累积的自定义训练循环的类似 PyTorch 的实现之间有什么区别（如果有的话）：

gradient_accumulation_steps = 5
for batch_idx, batch in enumerate(dataset):
  x_batch, y_true_batch = batch
  y_pred_batch = model(x_batch)

  loss = loss_fn(y_true_batch, y_pred_batch)
  loss.backward()

  if (batch_idx + 1) % gradient_accumulation_steps == 0: # (assumption: the number of batches is a multiple of gradient_accumulation_steps)
    optimizer.step()
    optimizer.zero_grad()

y_true_batches, y_pred_batches = [], []
gradient_accumulation_steps = 5
for batch_idx, batch in enumerate(dataset):
  x_batch, y_true_batch = batch
  y_pred_batch = model(x_batch)

  y_true_batches.append(y_true_batch)
  y_pred_batches.append(y_pred_batch)

  if (batch_idx + 1) % gradient_accumulation_steps == 0: # (assumption: the number of batches is a multiple of gradient_accumulation_steps)
    y_true = stack_vertically(y_true_batches)
    y_pred = stack_vertically(y_pred_batches)

    loss = loss_fn(y_true, y_pred)
    loss.backward()
    optimizer.step()
    optimizer.zero_grad()

    y_true_batches.clear()
    y_pred_batches.clear()

另外，作为一个不相关的问题：由于梯度累积的目的是在内存限制的情况下模拟更大的批量大小，这是否意味着我也应该按比例提高学习率？

Answer 1

1.两种方案的区别：
从概念上讲，您的两个实现是相同的：您为每次权重更新转发 gradient_accumulation_steps 批次。
正如您已经观察到的，第二种方法比第一种方法需要更多的内存资源。

但是，有一点不同：通常，损失函数实现使用mean 来减少批次的损失。当您使用梯度累积（第一个实现）时，您在每个小批量上使用 mean 减少，但在累积的 gradient_accumulation_steps 小批量上使用 sum。为了确保累积梯度实现与大批量实现相同，您需要非常小心地减少损失函数。在许多情况下，您需要将累积损失除以gradient_accumulation_steps。详细实现见this answer。

---

2。批量大小和学习率： 学习率和批量大小确实相关。当增加批量大小时，通常会降低学习率。
参见，例如：
Samuel L. Smith、Pieter-Jan Kindermans、Chris Ying、Quoc V. Le、Don't Decay the Learning Rate, Increase the Batch Size（ICLR 2018）。