注意
点击 这里 下载完整的示例代码
PyTorch 中清零梯度¶
创建于: 2020 年 4 月 20 日 | 最后更新于: 2024 年 6 月 13 日 | 最后验证于: 2024 年 11 月 5 日
在构建神经网络时，清零梯度是有益的。这是因为默认情况下，每当调用 .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. 导入加载数据所需的所有库¶
对于本代码示例，我们将仅使用 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 中清零梯度。
