Question

我正试图了解PyTorch中的梯度累积的内部工作原理。我的问题与这两个问题有关：

Why do we need to call zero_grad() in PyTorch?

Why do we need to explicitly call zero_grad()?

对第二个问题的已接受答案的评论表明，如果小批量太大而无法在单个前向通行中执行梯度更新，则可以使用累积的梯度，因此必须将其分为多个子批次。

考虑以下玩具示例：

import numpy as np
import torch


class ExampleLinear(torch.nn.Module):

    def __init__(self):
        super().__init__()
        # Initialize the weight at 1
        self.weight = torch.nn.Parameter(torch.Tensor([1]).float(),
                                         requires_grad=True)

    def forward(self, x):
        return self.weight * x


if __name__ == "__main__":
    # Example 1
    model = ExampleLinear()

    # Generate some data
    x = torch.from_numpy(np.array([4, 2])).float()
    y = 2 * x

    optimizer = torch.optim.SGD(model.parameters(), lr=0.01)

    y_hat = model(x)          # forward pass

    loss = (y - y_hat) ** 2
    loss = loss.mean()        # MSE loss

    loss.backward()           # backward pass

    optimizer.step()          # weight update

    print(model.weight.grad)  # tensor([-20.])
    print(model.weight)       # tensor([1.2000]

这正是人们所期望的结果。现在假设我们要使用梯度累积来逐样本处理数据集：

    # Example 2: MSE sample-by-sample
    model2 = ExampleLinear()
    optimizer = torch.optim.SGD(model2.parameters(), lr=0.01)

    # Compute loss sample-by-sample, then average it over all samples
    loss = []
    for k in range(len(y)):
        y_hat = model2(x[k])
        loss.append((y[k] - y_hat) ** 2)
    loss = sum(loss) / len(y)

    loss.backward()     # backward pass
    optimizer.step()    # weight update

    print(model2.weight.grad)  # tensor([-20.])
    print(model2.weight)       # tensor([1.2000]

与预期的一样，调用.backward()方法时将计算梯度。

最后，我的问题是：“幕后”到底发生了什么？

我的理解是，对于<PowBackward>变量，计算图是从<AddBackward>到<DivBackward> loss个操作动态更新的，并且没有关于用于除loss张量外，每个前向通过都保留在任何地方，张量可以更新直到后向通过。

以上段落中的推理是否有警告？最后，使用梯度累积时是否有最佳做法可循（例如，我在示例2 适得其反中使用的方法是否可以实现）？

Answer 1

您实际上并不是在累积渐变。如果您只有一个optimizer.zero_grad()通话，则仅放弃.backward()无效，因为渐变从一开始就已经为零（技术上为None，但它们将是自动初始化为零）。

两个版本之间的唯一区别是最终损失的计算方式。第二个示例的for循环执行与第一个示例中的PyTorch相同的计算，但是您必须分别进行计算，并且PyTorch无法优化（并行化和矢量化）您的for循环，这在GPU上产生了特别惊人的差异，前提是张量不是很小。

在进行梯度累积之前，让我们从您的问题开始：

最后一个问题是：“幕后”到底发生了什么？

当且仅当其中一个操作数已经是计算图的一部分时，才会在计算图中跟踪对张量的每个操作。设置张量的requires_grad=True时，它将创建具有单个顶点的计算图，即张量本身，该顶点将保留在图中。具有该张量的任何操作都会创建一个新的顶点，这是该操作的结果，因此从操作数到该顶点都有一条边，跟踪执行的操作。

a = torch.tensor(2.0, requires_grad=True)
b = torch.tensor(4.0)
c = a + b # => tensor(6., grad_fn=<AddBackward0>)

a.requires_grad # => True
a.is_leaf # => True

b.requires_grad # => False
b.is_leaf # => True

c.requires_grad # => True
c.is_leaf # => False

每个中间张量都自动需要梯度并具有grad_fn，该函数用于计算相对于其输入的偏导数。多亏了链式规则，我们可以以相反的顺序遍历整个图，以计算每个叶子的导数，这是我们要优化的参数。这就是反向传播的思想，也称为反向模式微分。有关更多详细信息，我建议阅读Calculus on Computational Graphs: Backpropagation。

PyTorch使用了确切的思路，当您调用loss.backward()时，它将以相反的顺序从loss开始遍历图形，并计算每个顶点的导数。每当到达叶子时，该张量的计算导数都存储在其.grad属性中。

在您的第一个示例中，这将导致：

MeanBackward -> PowBackward -> SubBackward -> MulBackward`

第二个示例几乎相同，除了您手动计算均值之外，对于损失计算的每个元素，您都没有多条路径，而是有一条单独的损失路径。要澄清的是，单一路径还可以计算每个元素的导数，但是在内部，这再次为某些优化打开了可能性。

# Example 1
loss = (y - y_hat) ** 2
# => tensor([16.,  4.], grad_fn=<PowBackward0>)

# Example 2
loss = []
for k in range(len(y)):
    y_hat = model2(x[k])
    loss.append((y[k] - y_hat) ** 2)
loss
# => [tensor([16.], grad_fn=<PowBackward0>), tensor([4.], grad_fn=<PowBackward0>)]

无论哪种情况，都会创建一个仅反向传播一次的图形，这就是不将其视为梯度累积的原因。

渐变累积

梯度累积是指在更新参数之前执行多次后退遍历的情况。目标是为多个输入（批次）具有相同的模型参数，然后基于所有这些批次更新模型的参数，而不是在每个批次之后执行更新。

让我们回顾一下您的例子。 x的大小为 [2] ，这就是我们整个数据集的大小。由于某些原因，我们需要基于整个数据集计算梯度。使用2的批处理大小自然是这种情况，因为我们可以一次拥有整个数据集。但是，如果我们只能批量生产1个尺寸，会发生什么？我们可以单独运行它们并像往常一样在每个批次之后更新模型，但是那样我们就不计算整个数据集的梯度。

我们需要做的是，使用相同的模型参数分别运行每个样本，并在不更新模型的情况下计算梯度。现在您可能在想，这不是您在第二版中所做的吗？几乎（但不是完全），您的版本中存在一个关键问题，即您使用的内存量与第一个版本相同，因为您具有相同的计算，因此计算图中的值数量也相同。

我们如何释放内存？我们需要摆脱前一批的张量以及计算图，因为这会占用大量内存来跟踪进行反向传播所需的所有信息。调用.backward()时，计算图会自动销毁（除非指定了retain_graph=True）。

def calculate_loss(x: torch.Tensor) -> torch.Tensor:
    y = 2 * x
    y_hat = model(x)
    loss = (y - y_hat) ** 2
    return loss.mean()


# With mulitple batches of size 1
batches = [torch.tensor([4.0]), torch.tensor([2.0])]

optimizer.zero_grad()
for i, batch in enumerate(batches):
    # The loss needs to be scaled, because the mean should be taken across the whole
    # dataset, which requires the loss to be divided by the number of batches.
    loss = calculate_loss(batch) / len(batches)
    loss.backward()
    print(f"Batch size 1 (batch {i}) - grad: {model.weight.grad}")
    print(f"Batch size 1 (batch {i}) - weight: {model.weight}")

# Updating the model only after all batches
optimizer.step()
print(f"Batch size 1 (final) - grad: {model.weight.grad}")
print(f"Batch size 1 (final) - weight: {model.weight}")

输出（为了便于阅读，我删除了包含的消息）：

Batch size 1 (batch 0) - grad: tensor([-16.])
Batch size 1 (batch 0) - weight: tensor([1.], requires_grad=True)
Batch size 1 (batch 1) - grad: tensor([-20.])
Batch size 1 (batch 1) - weight: tensor([1.], requires_grad=True)
Batch size 1 (final) - grad: tensor([-20.])
Batch size 1 (final) - weight: tensor([1.2000], requires_grad=True)

如您所见，在累积梯度的同时，所有批次的模型均保持相同的参数，最后只进行了一次更新。请注意，损失需要按批次进行缩放，以使整个数据集具有与使用单个批次相同的重要性。

虽然在此示例中，整个数据集在执行更新之前就已使用，但是您可以轻松地更改它，以在一定数量的批次后更新参数，但是您必须记住在采取优化步骤之后将梯度归零。一般配方为：

accumulation_steps = 10
for i, batch in enumerate(batches):
    # Scale the loss to the mean of the accumulated batch size
    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()

您可以在HuggingFace - Training Neural Nets on Larger Batches: Practical Tips for 1-GPU, Multi-GPU & Distributed setups中找到用于处理大批量的配方和更多技术。

了解PyTorch中的累积梯度

1 个答案:

渐变累积