4.学习率调度器

Dachao2026年8月7日大约 10 分钟

第 4 章 — 学习率调度器（lr_scheduler）

第 3 章把学习率从硬编码解放成了命令行参数 --lr，但整个训练过程它始终是一个固定值。本章解决这个问题：引入 torch.optim.lr_scheduler（学习率调度器），让学习率随训练进度动态变化——前期大步快跑，训练后期小步精调。

一、本章要解决的问题

固定学习率是一个两难：--lr 设大了，前期收敛快，但后期会在最优值附近来回震荡、loss 压不下去；设小了，后期精细，但前期像蜗牛爬，白白浪费训练时间。
上一章的训练日志里你会看到：val loss 到了 20 个 epoch 左右基本走平，出现一条"水平尾巴"——这不是模型能力到头了，而是固定步长不再适配。
判断标准：训练日志里出现一列 lr，且它随 epoch 稳步递减；对比"固定 lr"与"加了调度"两组实验，最终 test acc 有可观察的提升。

二、核心概念速览

下面 6 个概念是理解本章代码的全部"生词"。老规矩，先花 5 分钟读完，再看代码会轻松很多。

1. 学习率（lr）：下山时每步跨多大

把训练想象成摸黑下山：损失函数是一座山，参数是你脚下的坐标，目标是走到最低点。学习率就是每一步跨多大 ——跨得大，下山快，但容易一步冲过头；跨得小，走得稳，但龟速。第 3 章的 --lr 1e-3 说的就是这一步跨出 0.001 那么大。它是全训练里最敏感的超参数之一。

2. 为什么固定学习率后期收敛差：大步子在谷底来回震荡

刚下山时离谷底远，大跨步赶路没问题；可一旦接近谷底，如果步幅还是那么大，你会从谷底左边冲到右边，来回横跳，永远站不到正中央——这就是"震荡"。对应到训练里：固定 lr 下，loss 前期掉得飞快，后期却在一个平台上下抖动、怎么都压不下去。不是模型不会学了，而是"步子太大"，够不到更精细的最优位置。

3. 学习率调度：训练后期把步子变小

解法很朴素：前期大步快跑快速接近谷底，训练中后期逐步把步幅缩小，让参数稳稳落进谷底。这就是学习率调度（learning rate scheduling）。PyTorch 用 torch.optim.lr_scheduler 提供现成的"变速策略"，你选一种、告诉它节奏，它负责在每个 epoch 结束时自动改 lr，训练循环本身几乎不用动。

4. StepLR 与 CosineAnnealingLR：阶梯式折扣 vs 平滑下滑

StepLR 是阶梯式：每训练 N 轮（step_size），把当前 lr 直接乘一个折扣 gamma。比如 lr=1e-3、step_size=15、gamma=0.1，就是第 15、30 轮时 lr 依次变成 1e-4、1e-5，曲线像下楼梯——简单直白、可解释性强，但"隔多久折、折多狠"要自己试。CosineAnnealingLR 是平滑式：lr 沿一条余弦曲线从初始值缓缓降到 0，每个 epoch 都变一点，像滑滑梯。区别一句话：StepLR 隔一段突然打折，cosine 全程平滑下滑；cosine 免调 step_size/gamma 这类额外超参，实践中往往收敛更稳、效果更好。

5. scheduler.step() 与 optimizer.step()：一个按 epoch、一个按 batch

optimizer.step() 在训练循环里每个 batch 都调用一次，作用是"按这一批算出的梯度更新参数"；
scheduler.step() 则通常放在每个 epoch 结束后调用一次，作用是"按训练进度把 lr 调一档"。

类比：前者像每走一步调整落脚点，后者像每走一段路程调整"接下来每步跨多大"。千万别把 scheduler.step() 放进 batch 循环——lr 会被疯狂降掉，训练直接崩。

6. 为什么 cosine 的 T_max 要等于总 epoch 数

T_max 的意思是"用多少个 epoch 把 lr 从初始值降到 0"。如果 T_max 小于总 epoch 数，lr 提前归零，剩下的轮次相当于空跑不学习；如果大于，训练结束时 lr 还没降完，效果打折扣。让 T_max = --epochs，余弦曲线刚好在最后一个 epoch 落到 0，调度节奏与训练时长严丝合缝地匹配。

三、解决思路

引入 torch.optim.lr_scheduler，用新增参数 --lr-scheduler 在两种策略间切换：step（阶梯式）与 cosine（余弦式）。
构建时机：调度器要"挂"在优化器外面（把 optimizer 传进去），它只负责修改优化器内部的 lr，不碰参数的梯度更新。
调用时机：scheduler.step() 在每个 epoch 结束后调用一次，与按 batch 调用的 optimizer.step() 形成"参数微调 + 节奏变速"的双层节奏。
把当前 lr 打印进每轮日志（optimizer.param_groups[0]["lr"]），用肉眼验证调度曲线真的在走。

trade-off：调度器是 torch.optim 自带的模块，零新依赖，收益却可能很可观；代价是引入新概念（构建时机、step 时机、T_max）和几个新命令行参数，学习曲线多一点点。

四、代码变更

第 3 章引入的 tqdm 进度条本章继续沿用（Train/Val 两个循环的 tqdm(loader, desc=...) 包装不需要改动）。下面的 diff 只列出本章真正的新增/修改：

# 新增 3 个学习率调度参数 
  parser.add_argument("--lr", type=float, default=1e-3, help="学习率")
+ parser.add_argument("--lr-scheduler", type=str, default="step",
+                     choices=["step", "cosine"], help="学习率调度策略")
+ parser.add_argument("--lr-step-size", type=int, default=15, help="StepLR 每多少轮衰减一次")
+ parser.add_argument("--lr-gamma", type=float, default=0.1, help="StepLR 衰减系数")
  ...

# 在优化器之后构建调度器（挂在优化器"外面"）
  model = SimpleCNN().to(device)
  criterion = nn.CrossEntropyLoss()
  optimizer = optim.Adam(model.parameters(), lr=args.lr)
+
+ # 调度器：挂在优化器"外面"，按 epoch 修改 lr
+ if args.lr_scheduler == "step":
+     scheduler = optim.lr_scheduler.StepLR(
+         optimizer, step_size=args.lr_step_size, gamma=args.lr_gamma)
+ elif args.lr_scheduler == "cosine":
+     scheduler = optim.lr_scheduler.CosineAnnealingLR(
+         optimizer, T_max=args.epochs)   # T_max 必须等于总 epoch 数
+ else:
+     scheduler = None
  ...

# 主循环：每个 epoch 结束后推进一次调度器，并打印当前 lr 
  for epoch in range(1, args.epochs + 1):
      train_loss, train_acc = train_one_epoch(model, train_loader, criterion, optimizer, device)
      val_loss, val_acc = validate(model, val_loader, criterion, device)
+
+     # 关键：调度器在每个 epoch 结束后 step 一次（不同于优化器按 batch step）
+     if scheduler is not None:
+         scheduler.step()
+     current_lr = optimizer.param_groups[0]["lr"]
+
      metrics["train_loss"].append(train_loss)
      ...
      print(f"Epoch {epoch:02d}/{args.epochs} | "
            f"Train loss {train_loss:.4f} acc {train_acc:.4f} | "
-           f"Val   loss {val_loss:.4f} acc {val_acc:.4f}")
+           f"Val   loss {val_loss:.4f} acc {val_acc:.4f} | "
+           f"lr {current_lr:.2e}")

五、完整代码

创建 train.py（完整版，第 3 章基础上加入学习率调度）：

import argparse
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, random_split
from torchvision import datasets, transforms
from collections import defaultdict
from tqdm import tqdm   # 进度条（第 3 章引入，本章沿用）


def parse_args():
    parser = argparse.ArgumentParser(description="CIFAR-10 图像分类训练")
    # 训练超参数
    parser.add_argument("--epochs", type=int, default=10, help="训练轮数")
    parser.add_argument("--batch-size", type=int, default=64, help="每个 batch 的样本数")
    parser.add_argument("--lr", type=float, default=1e-3, help="初始学习率")
    # 本章新增：学习率调度相关参数
    parser.add_argument("--lr-scheduler", type=str, default="step",
                        choices=["step", "cosine"], help="学习率调度策略")
    parser.add_argument("--lr-step-size", type=int, default=15, help="StepLR 每多少轮衰减一次")
    parser.add_argument("--lr-gamma", type=float, default=0.1, help="StepLR 衰减系数")
    # 数据与运行环境
    parser.add_argument("--data-dir", type=str, default="./data", help="数据集存放目录")
    parser.add_argument("--num-workers", type=int, default=2, help="DataLoader 数据加载进程数")
    parser.add_argument("--device", type=str, default="auto", help="运行设备：auto/cuda/cpu")
    return parser.parse_args()


args = parse_args()

# 设备：auto 时自动检测 GPU
if args.device == "auto":
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
else:
    device = torch.device(args.device)
print(f"Using device: {device}")


class SimpleCNN(nn.Module):
    def __init__(self, num_classes=10):
        super().__init__()
        # 特征提取：3 个卷积块，把 32x32 的图逐级压到 4x4
        self.features = nn.Sequential(
            nn.Conv2d(3, 32, kernel_size=3, padding=1),
            nn.BatchNorm2d(32),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(2),
            nn.Conv2d(32, 64, kernel_size=3, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(2),
            nn.Conv2d(64, 128, kernel_size=3, padding=1),
            nn.BatchNorm2d(128),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(2),
        )
        # 分类头：全局平均池化 + 展平 + 全连接
        self.classifier = nn.Sequential(
            nn.AdaptiveAvgPool2d(1),
            nn.Flatten(),
            nn.Linear(128, num_classes),
        )

    def forward(self, x):
        return self.classifier(self.features(x))


transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2470, 0.2435, 0.2616)),
])

train_dataset = datasets.CIFAR10(root=args.data_dir, train=True, download=True, transform=transform)
test_dataset = datasets.CIFAR10(root=args.data_dir, train=False, download=True, transform=transform)
train_dataset, val_dataset = random_split(
    train_dataset, [45000, len(train_dataset) - 45000]
)

train_loader = DataLoader(train_dataset, batch_size=args.batch_size,
                          shuffle=True, num_workers=args.num_workers)
val_loader = DataLoader(val_dataset, batch_size=args.batch_size,
                        shuffle=False, num_workers=args.num_workers)
test_loader = DataLoader(test_dataset, batch_size=args.batch_size,
                         shuffle=False, num_workers=args.num_workers)

model = SimpleCNN().to(device)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=args.lr)

# 调度器：挂在优化器"外面"，按 epoch 修改 lr
if args.lr_scheduler == "step":
    scheduler = optim.lr_scheduler.StepLR(
        optimizer, step_size=args.lr_step_size, gamma=args.lr_gamma)
elif args.lr_scheduler == "cosine":
    scheduler = optim.lr_scheduler.CosineAnnealingLR(
        optimizer, T_max=args.epochs)   # T_max 必须等于总 epoch 数
else:
    scheduler = None


def train_one_epoch(model, loader, criterion, optimizer, device):
    model.train()
    total_loss, correct, total = 0.0, 0, 0
    for images, labels in tqdm(loader, desc="Train", leave=False):
        images, labels = images.to(device), labels.to(device)
        optimizer.zero_grad()
        outputs = model(images)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()   # 按 batch 更新一次参数（与调度器的 step 不同步）
        total_loss += loss.item() * images.size(0)
        correct += (outputs.argmax(dim=1) == labels).sum().item()
        total += images.size(0)
    return total_loss / total, correct / total


def validate(model, loader, criterion, device):
    model.eval()
    total_loss, correct, total = 0.0, 0, 0
    with torch.no_grad():
        for images, labels in tqdm(loader, desc="Val", leave=False):
            images, labels = images.to(device), labels.to(device)
            outputs = model(images)
            loss = criterion(outputs, labels)
            total_loss += loss.item() * images.size(0)
            correct += (outputs.argmax(dim=1) == labels).sum().item()
            total += images.size(0)
    return total_loss / total, correct / total


print("=" * 60)
print("训练配置：")
for k, v in vars(args).items():
    print(f"  {k:12s} = {v}")
print("=" * 60)

metrics = defaultdict(list)
for epoch in range(1, args.epochs + 1):
    train_loss, train_acc = train_one_epoch(model, train_loader, criterion, optimizer, device)
    val_loss, val_acc = validate(model, val_loader, criterion, device)

    # 关键：调度器在每个 epoch 结束后 step 一次（不同于优化器按 batch step）
    if scheduler is not None:
        scheduler.step()
    current_lr = optimizer.param_groups[0]["lr"]   # 读取当前学习率

    metrics["train_loss"].append(train_loss)
    metrics["train_acc"].append(train_acc)
    metrics["val_loss"].append(val_loss)
    metrics["val_acc"].append(val_acc)

    print(f"Epoch {epoch:02d}/{args.epochs} | "
          f"Train loss {train_loss:.4f} acc {train_acc:.4f} | "
          f"Val   loss {val_loss:.4f} acc {val_acc:.4f} | "
          f"lr {current_lr:.2e}")

test_loss, test_acc = validate(model, test_loader, criterion, device)
print(f"\n[Test] loss {test_loss:.4f} acc {test_acc:.4f}")

对比实验（各跑一条命令，重点看日志里的 lr 一列怎么变，以及最终 test acc）：

# StepLR：每 15 轮 lr 打 1 折（1e-3 -> 1e-4 -> 1e-5），曲线像下楼梯
python train.py --epochs 30 --lr-scheduler step --lr-step-size 15 --lr-gamma 0.1

# CosineAnnealingLR：lr 沿余弦曲线从 1e-3 平滑衰减到 0，曲线像滑梯
python train.py --epochs 30 --lr-scheduler cosine

六、本章小结

学到了什么
- 学习率是"下山步幅"：前期要大、后期要小；调度器就是自动控制步幅的节拍器。
- StepLR 阶梯式打折（step_size 间隔、gamma 折扣），CosineAnnealingLR 余弦平滑归零（T_max 取总 epoch 数）。
- 调度器挂在优化器外面：scheduler.step() 每个 epoch 调一次，optimizer.step() 每个 batch 调一次，节奏完全不同。
- 用 optimizer.param_groups[0]["lr"] 读取当前 lr 并打印进日志，是验证调度是否生效的最直观手段。
常见坑
- scheduler.step() 放错位置：放进 batch 循环会让 lr 骤降、训练崩溃；应放在每个 epoch 的验证完成之后。
- T_max 与 --epochs 不一致：cosine 曲线要么提前归零空跑、要么训练结束时还没降完。
- 中途改 --epochs：cosine 的剩余衰减路径会整体重算，不同 epoch 数的实验结果不能直接横向对比。
- 之后章节保存 checkpoint 时，要连 scheduler.state_dict() 一起存，否则恢复训练时调度状态会丢（第 5 章展开）。
下一章预告：训练到一半断电，几小时白跑；想回头选"验证集最优的那个 epoch"也没有存档——第 5 章引入 checkpoint 保存与恢复。

七、动手练习

画出调度曲线：用 step 和 cosine 各跑 30 epoch，把日志里的 lr 一列抄下来画折线图，对比"楼梯"与"滑梯"两种衰减形态。
对照实验：固定 --epochs 30，分别跑固定 lr（临时把调度器换成 None）、StepLR、cosine 三组，比较最终 test acc，体会"学习率衰减带来的收益"有多大。
加一个 none 选项：把 --lr-scheduler 的 choices 改成 ["step", "cosine", "none"]，让 else 分支返回 None，跑一跑验证"不调度"时日志里 lr 恒定不变。
挑战：想想为什么很多论文里的 best_lr 是"找出来的"而不是"猜出来的"——调度器之后，lr 的取值空间和训练时长绑在了一起，这个视角对第 6 章的调参很有用。