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

注意

跳转至页面底部 下载完整示例代码。

在 PyTorch 中将梯度归零#

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

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

简介#

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

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

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

设置#

由于我们将在本教程中训练数据，如果你正在使用可运行的笔记本，最好将运行环境切换到 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 库