Pytorch中loss.backward()和optimizer.step()的联系

技术背景

在深度学习训练过程中，反向传播（Back Propagation）和梯度下降（Gradient Descent）是两个核心步骤。在PyTorch中，loss.backward()用于执行反向传播，计算损失函数关于模型参数的梯度；optimizer.step()用于根据计算得到的梯度更新模型的参数。然而，两者之间并没有显式的调用关系，那么优化器是如何知道从哪里获取损失的梯度的呢？

实现步骤

1. 定义模型和优化器

在初始化优化器时，需要明确告诉它要更新哪些模型参数。例如：

import torch
import torch.optim as optim

# 定义模型参数
x = torch.tensor([1., 2.], requires_grad=True)
# 定义优化器
optim = torch.optim.SGD([x], lr=0.001)

2. 前向传播和损失计算

进行前向传播，计算模型的输出，并根据输出和真实标签计算损失。

# 前向传播
y = 100 * x
# 计算损失
loss = y.sum()

3. 反向传播

调用loss.backward()计算损失关于模型参数的梯度。

# 反向传播
loss.backward()

4. 参数更新

调用optimizer.step()根据计算得到的梯度更新模型参数。

# 参数更新
optim.step()

核心代码

import torch
import torch.optim as optim

# 我们的“模型”
x = torch.tensor([1., 2.], requires_grad=True)
y = 100 * x

# 计算损失
loss = y.sum()

# 计算损失关于参数的梯度
print(x.grad)     # None
loss.backward()      
print(x.grad)     # tensor([100., 100.])

# 修改参数，减去梯度
optim = torch.optim.SGD([x], lr=0.001)
print(x)        # tensor([1., 2.], requires_grad=True)
optim.step()
print(x)        # tensor([0.9000, 1.9000], requires_grad=True)

最佳实践

梯度清零

在每个训练步骤中，需要调用optimizer.zero_grad()来清零梯度，因为在每个训练步骤中我们希望计算新的梯度，不清零梯度会导致梯度在不同批次之间累积。

# 在每个训练步骤开始时清零梯度
optim.zero_grad()

确保模型设置完成后再构建优化器

在初始化优化器后，如果将模型移动到GPU上，可能会导致优化器引用的参数出现问题。因此，要确保在构建优化器之前完成模型的设置。

常见问题

梯度未更新

如果在调用optimizer.step()后模型参数没有更新，可能是因为：

• 没有调用loss.backward()计算梯度。
• 参数的requires_grad属性设置为False。
• 梯度被意外清零或未正确计算。

梯度累积

如果没有在每个训练步骤中调用optimizer.zero_grad()，梯度会在不同批次之间累积，导致参数更新异常。因此，一定要在每个训练步骤开始时清零梯度。