评价此页
★ ★ ★ ★ ★
recipes/recipes/zeroing_out_gradients

注意

转到底部 下载完整的示例代码。

PyTorch 中归零梯度#

创建于：2020年4月20日 | 最后更新：2025年4月28日 | 最后验证：2024年11月05日

在构建神经网络时，归零梯度是有益的。这是因为默认情况下，每当调用 .backward() 时，梯度都会在缓冲区中累积（即不覆盖）。

简介#

在训练神经网络时，模型能够通过梯度下降来提高其准确性。简而言之，梯度下降是通过调整模型中的权重和偏差来最小化损失（或误差）的过程。

torch.Tensor 是 PyTorch 的核心类。当你创建一个张量时，如果将其属性 .requires_grad 设置为 True，则该包会跟踪其上的所有操作。这会在后续的反向传播中发生。该张量的梯度将累积到 .grad 属性中。当在损失张量上调用 .backward() 时，会计算所有梯度的累积（或总和）。

在某些情况下，可能需要归零张量的梯度。例如：当你开始训练循环时，你应该归零梯度，以便能够正确地进行此跟踪。在本教程中，我们将学习如何使用 PyTorch 库归零梯度。我们将通过在 PyTorch 内置的 CIFAR10 数据集上训练神经网络来演示这一点。

设置#

由于本教程将训练数据，如果你在一个可运行的 notebook 中，最好将运行时切换到 GPU 或 TPU。在开始之前，如果尚未安装 torch 和 torchvision，则需要进行安装。

pip install torchvision

步骤#

步骤 1 到 4 设置了我们的数据和用于训练的神经网络。归零梯度的过程发生在步骤 5 中。如果你已经构建了数据和神经网络，请跳到 5。

  1. 导入加载数据所需的所有库

  2. 加载和归一化数据集

  3. 构建神经网络

  4. 定义损失函数

  5. 在训练网络时归零梯度

1. 导入加载数据所需的库#

在本教程中，我们将仅使用 torch 和 torchvision 来访问数据集。

import torch

import torch.nn as nn
import torch.nn.functional as F

import torch.optim as optim

import torchvision
import torchvision.transforms as transforms

2. 加载和归一化数据集#

PyTorch 提供了各种内置数据集（有关更多信息，请参阅加载数据教程）。

transform = transforms.Compose(
    [transforms.ToTensor(),
     transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])

trainset = torchvision.datasets.CIFAR10(root='./data', train=True,
                                        download=True, transform=transform)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=4,
                                          shuffle=True, num_workers=2)

testset = torchvision.datasets.CIFAR10(root='./data', train=False,
                                       download=True, transform=transform)
testloader = torch.utils.data.DataLoader(testset, batch_size=4,
                                         shuffle=False, num_workers=2)

classes = ('plane', 'car', 'bird', 'cat',
           'deer', 'dog', 'frog', 'horse', 'ship', 'truck')

3. 构建神经网络#

我们将使用卷积神经网络。要了解更多信息，请参阅定义神经网络教程。

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(3, 6, 5)
        self.pool = nn.MaxPool2d(2, 2)
        self.conv2 = nn.Conv2d(6, 16, 5)
        self.fc1 = nn.Linear(16 * 5 * 5, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)

    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        x = x.view(-1, 16 * 5 * 5)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x

4. 定义损失函数和优化器#

让我们使用分类交叉熵损失和带动量的 SGD。

net = Net()
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)

5. 在训练网络时归零梯度#

这时事情开始变得有趣了。我们只需遍历数据迭代器，将输入馈送到网络并进行优化。

请注意，对于每个数据实体，我们都会归零梯度。这是为了确保我们在训练神经网络时不会跟踪任何不必要的信息。

for epoch in range(2):  # loop over the dataset multiple times

    running_loss = 0.0
    for i, data in enumerate(trainloader, 0):
        # get the inputs; data is a list of [inputs, labels]
        inputs, labels = data

        # zero the parameter gradients
        optimizer.zero_grad()

        # forward + backward + optimize
        outputs = net(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

        # print statistics
        running_loss += loss.item()
        if i % 2000 == 1999:    # print every 2000 mini-batches
            print('[%d, %5d] loss: %.3f' %
                  (epoch + 1, i + 1, running_loss / 2000))
            running_loss = 0.0

print('Finished Training')

你也可以使用 model.zero_grad()。只要你的所有模型参数都在该优化器中，它就与使用 optimizer.zero_grad() 相同。请自行判断使用哪种方法。

恭喜！你已成功在 PyTorch 中归零了梯度。

了解更多#

查看这些其他秘籍以继续您的学习

在此页面上
PyTorch 库