关于梯度累积的说明

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。有关详细实施，请参阅。

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

关于梯度累积的说明

Clarification about Gradient Accumulation

python

gradient-descent

pytorch