PyTorch中SGD优化器批量循环内固定LR的实现问题

1. 调整学习率调度器的更新时机 #

问题根源是你把scheduler.step()放在了batch循环内部，导致每个batch都更新学习率。正确做法是将scheduler.step()移到每个epoch循环的末尾，这样整个epoch内所有batch共享同一学习率，仅在epoch结束后更新下一轮的LR。

2. 实现最优模型保存逻辑 #

通过跟踪验证集指标（损失或准确率），实时保存表现最优的模型权重。

完整代码示例 #

import torch
import torch.nn as nn
from torch.optim import SGD
from torch.optim.lr_scheduler import LinearLR

# 示例模型（根据你的任务替换为实际模型）
model = nn.Linear(X_train.shape[1], 1)
# 损失函数（分类任务替换为CrossEntropyLoss）
criterion = nn.MSELoss()
optimizer = SGD(model.parameters(), lr=0.1)
# 初始化LinearLR：比如5个epoch内从初始LR降到0.1倍
scheduler = LinearLR(optimizer, start_factor=1.0, end_factor=0.1, total_iters=5)

# 初始化最优指标与模型存储路径
best_val_loss = float('inf')
best_model_path = 'best_model.pth'

num_epochs = 10
for epoch in range(num_epochs):
    # 训练模式
    model.train()
    train_loss = 0.0

    # Batch循环（这里按32样本分batch，替换为你的DataLoader或遍历方式）
    for X_batch, y_batch in zip(X_train.split(32), y_train.split(32)):
        optimizer.zero_grad()
        outputs = model(X_batch)
        loss = criterion(outputs, y_batch)
        loss.backward()
        optimizer.step()  # 仅更新模型参数，不改动LR
        train_loss += loss.item() * X_batch.size(0)
    
    train_loss /= len(X_train)

    # 验证阶段
    model.eval()
    val_loss = 0.0
    with torch.no_grad():
        for X_batch, y_batch in zip(X_test.split(32), y_test.split(32)):
            outputs = model(X_batch)
            loss = criterion(outputs, y_batch)
            val_loss += loss.item() * X_batch.size(0)
    val_loss /= len(X_test)

    print(f'Epoch {epoch+1}/{num_epochs}, 训练损失: {train_loss:.4f}, 验证损失: {val_loss:.4f}, 当前LR: {optimizer.param_groups[0]["lr"]:.6f}')

    # 保存最优模型
    if val_loss < best_val_loss:
        best_val_loss = val_loss
        torch.save(model.state_dict(), best_model_path)
        print(f'更新最优模型，当前验证损失: {best_val_loss:.4f}')

    # 每个epoch结束后更新学习率
    scheduler.step()

# 训练完成后加载最优模型
model.load_state_dict(torch.load(best_model_path))

关键说明 #

• Batch内LR固定：scheduler.step()仅在epoch末尾执行，确保整个epoch的所有batch使用同一学习率，避免LR过早下降。
• 最优模型保存：通过比较验证集指标（这里用损失，分类任务可替换为准确率），保留性能最好的模型权重，规避训练后期过拟合导致的性能下滑。

内容的提问来源于stack exchange，提问作者Marios Karaoulis