optimizer.zero_grad()解释

作用

在每次 反向传播（backward） 之前，将优化器管理的所有参数的梯度清零。这一步非常重要，因为 PyTorch 默认会 累加 梯度，而不是自动覆盖。通过清零，可以确保梯度计算仅基于当前的前向和后向传播。

为什么需要清零梯度？

在 PyTorch 中，梯度计算遵循如下流程：

1. 前向传播：计算当前输入的模型输出。
2. 损失计算：根据预测输出和真实标签计算损失。
3. 反向传播：通过 loss.backward() 计算损失相对于模型参数的梯度。
4. 参数更新：通过 optimizer.step() 根据梯度更新模型参数。

但是，PyTorch 的梯度是默认累加的。如果不清零梯度，前一次反向传播计算的梯度值会叠加到下一次中，导致错误的优化结果。例如：

# 模拟梯度累加问题
import torch
from torch import nn, optim

model = nn.Linear(2, 1)  # 简单线性模型
optimizer = optim.SGD(model.parameters(), lr=0.01)

# 假设两次输入
input1 = torch.tensor([[1.0, 2.0]])
input2 = torch.tensor([[3.0, 4.0]])
target = torch.tensor([[1.0]])

criterion = nn.MSELoss()

# 第一次前向传播和反向传播
output1 = model(input1)
loss1 = criterion(output1, target)
loss1.backward()  # 梯度计算

# 如果不清零，第二次的梯度将叠加
output2 = model(input2)
loss2 = criterion(output2, target)
loss2.backward()

# 梯度累加，导致错误
print(model.weight.grad)  # 包含两次梯度的累加

结果中梯度被累加，这不是我们希望的。因此，需要在每次反向传播前调用 optimizer.zero_grad() 清零梯度。

使用方法

一般 optimizer.zero_grad() 和 loss.backward() 成对出现，常见流程如下：

for epoch in range(num_epochs):
    for data, target in dataloader:
        # 前向传播
        output = model(data)
        loss = criterion(output, target)

        # 清零梯度
        optimizer.zero_grad()

        # 反向传播
        loss.backward()

        # 更新参数
        optimizer.step()

为什么梯度会累加？

PyTorch 为了方便处理某些特殊情况（如梯度累积用于大批量训练）默认梯度不会自动清零。此特性在以下情况下有用：

1. 小批量梯度累积：

   • 当 GPU 显存不足时，可以分多步计算梯度，并在多次累积后更新参数。

   • 例如：

     for i, (data, target) in enumerate(dataloader):
         output = model(data)
         loss = criterion(output, target)
         
         # 累积梯度，不清零
         loss.backward()
         
         if (i + 1) % accumulation_steps == 0:  # 每 N 次更新参数
             optimizer.step()
             optimizer.zero_grad()
     

2. 手动操作梯度：累积梯度允许手动调整或控制。

总结

optimizer.zero_grad() 的必要性：防止梯度累加，确保每次训练只基于当前批次的数据。

一般流程：

• 清零梯度 (optimizer.zero_grad())
• 反向传播 (loss.backward())
• 参数更新 (optimizer.step())

灵活性：如果有梯度累积需求，可以选择性地省略这一步。