为什么我们需要在PyTorch中调用zero_grad()？

Question

zero_grad()需要在训练期间调用该方法.但文档不是很有帮助

|  zero_grad(self)
|      Sets gradients of all model parameters to zero.

为什么我们需要调用这个方法？

Answer 1

在PyTorch,我们需要在开始进行反向传播之前将梯度设置为零,因为PyTorch 会在后续的后向传递中累积渐变.在训练RNN时这很方便.因此,默认操作是在每次loss.backward()调用时累积渐变.

因此,当您开始训练循环时,理想情况下应该zero out the gradients正确执行参数更新.否则,梯度将指向除预期方向之外的某些其他方向朝向最小值(或在最大化目标的情况下最大值).

这是一个简单的例子:

import torch
from torch.autograd import Variable
import torch.optim as optim

def linear_model(x, W, b):
    return torch.matmul(x, W) + b

data, targets = ...

W = Variable(torch.randn(4, 3), requires_grad=True)
b = Variable(torch.randn(3), requires_grad=True)

optimizer = optim.Adam([W, b])

for sample, target in zip(data, targets):
    # clear out the gradients of all Variables 
    # in this optimizer (i.e. W, b)
    optimizer.zero_grad()
    output = linear_model(sample, W, b)
    loss = (output - target) ** 2
    loss.backward()
    optimizer.step()

或者,如果你正在做一个香草梯度下降

W = Variable(torch.randn(4, 3), requires_grad=True)
b = Variable(torch.randn(3), requires_grad=True)

for sample, target in zip(data, targets):
    # clear out the gradients of Variables 
    # (i.e. W, b)
    W.grad.data.zero_()
    b.grad.data.zero_()

    output = linear_model(sample, W, b)
    loss = (output - target) ** 2
    loss.backward()

    W -= learning_rate * W.grad.data
    b -= learning_rate * b.grad.data

Answer 2

虽然这个想法可以从所选的答案中得出，但我觉得我想明确地写出来。

能够决定何时调用，optimizer.zero_grad()并optimizer.step()为优化器在训练循环中如何累积和应用梯度提供了更多自由。当模型或输入数据很大且 GPU 无法容纳一批训练时，这一点至关重要。

在此示例中，有两个参数，名为train_batch_size和gradient_accumulation_steps。

• train_batch_size是前向传递的批量大小，紧随loss.backward(). 这受到 GPU 内存的限制。

• gradient_accumulation_steps是实际的训练批量大小，其中累积了多次前向传递的损失。这不受GPU 内存的限制。

从这个例子中，您可以看到optimizer.zero_grad()may 后面如何跟上optimizer.step()but NOT loss.backward()。loss.backward()在每次迭代中都会被调用（第 216 行），但optimizer.zero_grad()仅optimizer.step()当累积的训练批次数量等于gradient_accumulation_steps（第 219 行块内的第 227 行if）时才会调用 和 。

另外，有人在询问 TensorFlow 中的等效方法。我猜tf.GradientTape具有相同的目的。

Answer 3

为什么要渐变？

梯度向优化器建议步入的方向。每次使用 处理一批输入时.backward()，都会积累走向何处的“建议”。请注意，建议比决定弱得多。当您调用 时optimizer.step()，优化器会使用这些建议来实际决定实际执行的步骤。这些决策可能会受到学习率、过去的步骤（例如动量）和过去的权重（例如SWA）的影响。优化器会读取建议，然后朝着希望能够最小化未来损失的方向迈进。

loss.backward()        # Compute gradients.
optimizer.step()       # Tell the optimizer the gradients, then step.
optimizer.zero_grad()  # Zero the gradients to start fresh next time.

为什么要将梯度归零？

完成一个步骤后，您实际上不需要跟踪之前的步骤建议（即梯度）。通过将梯度归零，您将丢弃这些信息。一些优化器已经在内部自动跟踪这些信息。

通过下一批输入，您可以从头开始建议下一步该做什么。这个建议是纯粹的，没有受到过去的影响。然后，您将这些“纯”信息提供给优化器，然后由优化器准确决定要执行的步骤。

当然，您可以决定保留以前的梯度，但该信息有些过时，因为您处于损失表面上的全新位置。谁敢说接下来最好的方向还是和以前一样？可能会完全不一样！这就是为什么最流行的优化算法会丢弃大部分过时的信息（通过将梯度归零）。

另一种选择：完全删除梯度（而不是归零）

您也可以完全删除它们，而不是将渐变归零。PyTorch性能调优指南建议：

loss.backward()        # Compute gradients.
optimizer.step()       # Tell the optimizer the gradients, then step.
optimizer.zero_grad()  # Zero the gradients to start fresh next time.

# CONSIDER:
for param in model.parameters():
    param.grad = None

...但其中一位开发人员在 5 年前的评论中提到了这一点：

主要区别在于，包含梯度的张量不会在每次向后传递时重新分配。由于内存分配非常昂贵（尤其是在 GPU 上），因此效率更高。

两者之间还存在其他细微差别，例如某些优化器在梯度为 0 或 None 时表现不同。确信还有其他地方也有类似的行为。

...另一方面，就地操作通常被认为是不必要的，甚至在某些情况下不是最佳的，所以我想 YMMV 与这两种方法的性能有关。