浅谈pytorch中为什么要用 zero_grad() 将梯度清零

下面是详细讲解pytorch中为什么要用zero_grad()将梯度清零的攻略。

什么是pytorch中的梯度？

在深度学习中，我们通常使用反向传播算法来计算模型的梯度。在pytorch中，模型的梯度保存在参数的grad属性中。

例如，以下代码创建了一个简单的网络，并计算了模型参数的梯度。

import torch
import torch.nn as nn

# 创建网络
net = nn.Linear(10, 1)

# 定义输入和目标
inputs = torch.randn(1, 10)
targets = torch.randn(1, 1)

# 计算损失
outputs = net(inputs)
loss = torch.nn.functional.mse_loss(outputs, targets)

# 计算梯度
loss.backward()

在这个例子中，对loss调用backward()方法会自动计算模型中所有参数的梯度，并将其保存在相应的参数的grad属性中。

为什么要使用zero_grad()将梯度清零？

在训练过程中，每次反向传播之后，模型的梯度会累加到之前的梯度上。当我们想要训练一个新的batch数据时，如果不清空已有的梯度，则这些梯度会对新的batch数据产生不必要的影响，从而影响到模型的训练效果。

例如，以下代码演示了在不清空梯度的情况下，连续进行两次反向传播的影响。

import torch
import torch.nn as nn

# 创建网络
net = nn.Linear(10, 1)

# 定义输入和目标
inputs1 = torch.randn(1, 10)
inputs2 = torch.randn(1, 10)
targets = torch.randn(1, 1)

# 计算损失1
outputs1 = net(inputs1)
loss1 = torch.nn.functional.mse_loss(outputs1, targets)

# 反向传播1
loss1.backward()

# 计算损失2
outputs2 = net(inputs2)
loss2 = torch.nn.functional.mse_loss(outputs2, targets)

# 反向传播2
loss2.backward()

# 打印参数的梯度
print(net.weight.grad)

在这个例子中，我们首先计算了一个损失loss1，进行一次反向传播，并将模型参数的梯度保存在grad属性中。然后，我们计算了另一个损失loss2，并进行一次反向传播。由于在第一次反向传播后我们没有清空模型参数的梯度，因此第二次反向传播计算的梯度会与第一次的梯度进行累加。最终，参数的梯度包含了两次损失的影响，导致模型训练结果产生错误。

为了避免这种情况的发生，我们需要在每次训练batch数据之前，使用zero_grad()方法将参数的梯度清零，以确保每个batch数据计算的梯度只包含自己的影响。

如何正确使用zero_grad()方法

在pytorch中，zero_grad()方法可以应用于网络中的所有参数。以下是一些示例代码，演示了如何正确使用这个方法。

import torch
import torch.nn as nn

# 创建网络
net = nn.Linear(10, 1)

# 定义输入和目标
inputs1 = torch.randn(1, 10)
inputs2 = torch.randn(1, 10)
targets = torch.randn(1, 1)

# 创建优化器
optimizer = torch.optim.SGD(net.parameters(), lr=0.01)

# 计算损失1
outputs1 = net(inputs1)
loss1 = torch.nn.functional.mse_loss(outputs1, targets)

# 反向传播1
optimizer.zero_grad()     # 清零梯度
loss1.backward()
optimizer.step()          # 更新参数

# 计算损失2
outputs2 = net(inputs2)
loss2 = torch.nn.functional.mse_loss(outputs2, targets)

# 反向传播2
optimizer.zero_grad()     # 清零梯度
loss2.backward()
optimizer.step()           # 更新参数

# 打印参数的梯度
print(net.weight.grad)

在这个例子中，我们使用了SGD优化器进行参数更新。在每个batch数据训练之前，我们首先使用zero_grad()方法将模型参数的梯度清零。然后，我们计算了第一个batch数据的损失loss1，进行一次反向传播，并使用优化器更新了参数。接下来，我们计算了第二个batch数据的损失loss2，并进行了一次反向传播和参数更新。在这个过程中，我们使用zero_grad()方法在每次训练batch数据之前清空了参数梯度，确保每个batch数据的梯度只包含自己的影响。

综上所述，使用zero_grad()方法可以确保每个batch数据计算的梯度只包含自己的影响，从而保证模型训练的正确性。