训练时忽略优化器参数更新？梯度未更新的潜在原因与修复方案

# optimizer_update_bug.py
import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
​
torch.manual_seed(0)
​
class SimpleMLP(nn.Module):
    def __init__(self):
        super(SimpleMLP, self).__init__()
        self.fc1 = nn.Linear(28 * 28, 128)
        self.fc2 = nn.Linear(128, 10)
        
    def forward(self, x):
        x = torch.flatten(x, 1)
        x = F.relu(self.fc1(x))
        x = self.fc2(x)
        return x
​
def make_loader(n=1024, batch_size=64):
    X = torch.randn(n, 28, 28)
    y = torch.randint(0, 10, (n,))
    dataset = torch.utils.data.TensorDataset(X, y)
    return torch.utils.data.DataLoader(dataset, batch_size=batch_size, shuffle=True)
​
def train(model, dataloader, optimizer, epochs=10):
    for epoch in range(epochs):
        model.train()
        total_loss = 0
        correct = 0
        total = 0
        for data, target in dataloader:
            optimizer.zero_grad()  # 清空梯度
            output = model(data)  # 前向传播
            loss = F.cross_entropy(output, target)  # 计算损失
            loss.backward()  # 反向传播
            
            optimizer.step()  # 更新参数
            total_loss += loss.item()
​
            _, predicted = output.max(1)
            correct += (predicted == target).sum().item()
            total += target.size(0)
​
        avg_loss = total_loss / len(dataloader)
        accuracy = 100. * correct / total
        print(f"Epoch {epoch+1}, Loss: {avg_loss:.4f}, Accuracy: {accuracy:.2f}%")
​
def experiment():
    model = SimpleMLP()
    optimizer = optim.SGD(model.parameters(), lr=0.1)
    dataloader = make_loader()
​
    # 错误的优化器参数更新
    optimizer.zero_grad()  # 此处应该被注释，避免清空梯度后未进行更新
    for epoch in range(5):
        train(model, dataloader, optimizer)  # 训练 5 轮
    print("Training completed.")
​
if __name__ == "__main__":
    experiment()

optimizer.zero_grad()  # 确保每次开始计算新的梯度时，梯度清零

optimizer.step()  # 更新参数

for param in model.parameters():
    if param.grad is None:
        print(f"Parameter {param} has no gradient.")

optimizer = optim.SGD(model.parameters(), lr=0.1)  # 调整学习率

for epoch in range(epochs):
    model.train()
    for data, target in dataloader:
        optimizer.zero_grad()  # 确保每次梯度计算前清空旧梯度
        output = model(data)
        loss = F.cross_entropy(output, target)
        loss.backward()
        optimizer.step()  # 更新参数

# 修复训练过程中的梯度更新问题
optimizer.zero_grad()  # 清空梯度
loss.backward()        # 反向传播
optimizer.step()       # 更新模型参数

for param in model.parameters():
    if param.grad is None:
        print(f"Parameter {param} has no gradient.")

for param_group in optimizer.param_groups:
    print("Current learning rate:", param_group['lr'])

import torch
print(f"Memory allocated: {torch.cuda.memory_allocated()} bytes")