python - 关于梯度累积的说明

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)。