fix(trainer): 检测训练续训以配置为准，并恢复学习率调度进度

train.py --resume 走的是与 SimCLR 不同的一套代码，同样的问题在这里更严重：

- load_checkpoint 直接 optimizer.load_state_dict，把 lr/weight_decay 连同动量缓冲
  一起还原成存档值：yaml 写 lr=5e-4，实际仍按检查点的 1e-3 训
- save_checkpoint 根本没保存调度器状态，续训时 scheduler.last_epoch 从 0 开始，
  而 optimizer 的 lr 又来自检查点，两者互相矛盾，cosine 曲线完全错乱

改为与 SimCLR 一致的「配置文件是唯一真相源」：

- save_checkpoint 补上 scheduler 状态、global_step 与当时的 train 配置快照
- 新增 Trainer.resume()：加载优化器状态后按组写回配置里的超参（只留动量缓冲，
  并保住 BN/bias 那组 weight_decay=0 的结构），调度器按当前配置重建再快进到已训 epoch
- 配置与检查点不一致时逐项打日志说明按当前配置执行
- train.py 改用 trainer.resume()，不再自己拼恢复逻辑

实测续训到 10 轮时 lr=3.273e-04，与 cosine 在 4/10 处的理论值精确一致；
weight_decay、T_max、调度进度、动量缓冲均正确。

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

README.md

@@ -109,6 +109,19 @@ uv run python scripts/export_onnx.py --config configs/train/resnet18.yaml \
 
 产物写在 `work_dir`（默认 `runs/<name>/`）：`train.log`、`config.yaml`、`last.pth`、`best.pth`、`swanlog/`。
 
+**续训以配置为准。** `--resume` 只从检查点取"进度"——模型权重、优化器动量缓冲、
+已训到第几轮、`best_metric`、`global_step`；`epochs`、`lr`、`weight_decay`、调度器
+一律以当前 yaml 为准，所以可以顺手调超参：
+
+```bash
+uv run python train.py --config configs/train/baseline.yaml --resume runs/baseline/last.pth \
+    --set train.epochs=200 train.optimizer.lr=5e-4
+```
+
+与检查点不一致时日志会逐项列出（`配置与检查点不一致，按当前配置执行：总轮数 100 -> 200；
+optimizer.lr 0.001 -> 0.0005`）。学习率调度按新的总轮数重建并快进到当前 epoch。
+`epochs` 是**总轮数**不是剩余轮数。SimCLR 自监督那边是同样的语义。
+
 ## 多类 / 单类分类
 
 图像级分类的形式由 yaml 的 `label_mode` 决定，模型的分类头、分类损失、分类指标都跟着切换

engine/checkpoint.py

@@ -16,16 +16,27 @@ def save_checkpoint(
     epoch: int = 0,
     best_metric: float | None = None,
     config: dict | None = None,
+    scheduler: Any = None,
+    global_step: int = 0,
+    train_cfg: dict | None = None,
 ) -> Path:
+    """保存检查点。
+
+    除模型与优化器外还记下调度器状态、global_step 和当时的 train 配置：
+    少了这些，续训时学习率调度会从头开始，而且无法判断 yaml 里的超参有没有被改过。
+    """
     path = Path(path)
     path.parent.mkdir(parents=True, exist_ok=True)
     torch.save(
         {
             "model": model.state_dict(),
             "optimizer": optimizer.state_dict() if optimizer else None,
+            "scheduler": scheduler.state_dict() if scheduler else None,
             "epoch": epoch,
+            "global_step": global_step,
             "best_metric": best_metric,
             "config": config,
+            "train_cfg": train_cfg,
         },
         path,
     )

engine/trainer.py

@@ -4,6 +4,7 @@ from __future__ import annotations
 
 import math
 import time
+import warnings
 from collections.abc import Mapping
 from pathlib import Path
 from typing import Any
@@ -12,7 +13,7 @@ import torch
 from torch import nn
 from torch.utils.data import DataLoader
 
-from engine.checkpoint import save_checkpoint
+from engine.checkpoint import load_checkpoint, save_checkpoint
 from engine.evaluator import Evaluator, _to_device
 from engine.losses import DetectionLoss
 from engine.metrics import AverageMeter
@@ -149,6 +150,87 @@ class Trainer:
         self.peak_memory.reset()
         return stats
 
+    def resume(self, path: str | Path) -> dict:
+        """从检查点续训。
+
+        **配置文件是唯一真相源**：检查点只提供"进度"——模型权重、优化器动量缓冲、
+        已训到第几个 epoch、global_step、best_metric；而 lr、weight_decay、总轮数、
+        调度器一律以当前 yaml 为准。否则改了配置再续训，改动会被检查点静默覆盖。
+        """
+        state = load_checkpoint(path, self.model, map_location=self.device)
+
+        if state.get("optimizer"):
+            # optimizer.load_state_dict 会整体覆盖 param_groups，把 lr/weight_decay
+            # 一并还原成存档时的值。先抄下配置构造出的超参，加载完按组写回，
+            # 只留下动量缓冲。按组写回也保住了 BN 与 bias 那组 weight_decay=0 的结构。
+            intended = [
+                {key: value for key, value in group.items() if key != "params"}
+                for group in self.optimizer.param_groups
+            ]
+            self.optimizer.load_state_dict(state["optimizer"])
+            for group, wanted in zip(self.optimizer.param_groups, intended, strict=True):
+                group.update(wanted)
+
+        self.start_epoch = int(state.get("epoch", -1)) + 1
+        self.global_step = int(state.get("global_step", 0))
+        best = state.get("best_metric")
+        self.best_metric = float(best) if best is not None and math.isfinite(float(best)) else float("-inf")
+
+        self._restore_scheduler()
+        self._log_config_changes(state)
+        logger.info(
+            "已从 %s 恢复：起始 epoch=%d，best_%s=%.4f",
+            path,
+            self.start_epoch,
+            self.monitor,
+            self.best_metric,
+        )
+        return state
+
+    def _restore_scheduler(self) -> None:
+        """按**当前配置**重建学习率调度，再快进到已训到的 epoch。
+
+        不恢复检查点里的调度器状态：它带着当时的 T_max 与 base_lrs，一旦 yaml 里改了
+        总轮数或学习率就会被静默覆盖。cosine / step / multistep 都是 epoch 的确定性函数，
+        快进与恢复状态等价，所以统一走重建这一条路。
+        """
+        if self.scheduler is None:
+            return
+        # cosine 的链式实现基于当前 lr 递推，必须先复位回 initial_lr，
+        # 否则从检查点带来的 lr 出发会算出一条错的曲线。
+        for group in self.optimizer.param_groups:
+            group["lr"] = group.get("initial_lr", group["lr"])
+        self.scheduler = build_scheduler(self.optimizer, self.cfg.get("scheduler"), self.epochs)
+        with warnings.catch_warnings():
+            warnings.simplefilter("ignore")  # 快进阶段的 "step before optimizer.step()" 属预期
+            for _ in range(self.start_epoch):
+                self.scheduler.step()
+
+    def _log_config_changes(self, state: dict) -> None:
+        """检查点里的配置与当前配置不一致时说清楚用的是哪一份。"""
+        archived = state.get("train_cfg") or {}
+        changes = []
+        if archived.get("epochs") not in (None, self.epochs):
+            changes.append(f"总轮数 {archived['epochs']} -> {self.epochs}")
+        for key, value in (archived.get("optimizer") or {}).items():
+            current = dict(self.cfg.get("optimizer", {})).get(key)
+            if current is not None and current != value:
+                changes.append(f"optimizer.{key} {value} -> {current}")
+        archived_scheduler = (archived.get("scheduler") or {}).get("name")
+        current_scheduler = (dict(self.cfg.get("scheduler") or {})).get("name")
+        if archived_scheduler and archived_scheduler != current_scheduler:
+            changes.append(f"scheduler {archived_scheduler} -> {current_scheduler}")
+
+        if changes:
+            logger.warning("配置与检查点不一致，按**当前配置**执行：%s", "；".join(changes))
+        logger.info(
+            "续训超参：lr=%.2e（initial_lr=%.2e）weight_decay=%s，共 %d 轮",
+            self.optimizer.param_groups[0]["lr"],
+            self.optimizer.param_groups[0].get("initial_lr", float("nan")),
+            [g["weight_decay"] for g in self.optimizer.param_groups],
+            self.epochs,
+        )
+
     def fit(self) -> float:
         logger.info("开始训练：共 %d 个 epoch，输出目录 %s", self.epochs, self.work_dir)
         for epoch in range(self.start_epoch, self.epochs):
@@ -182,6 +264,9 @@ class Trainer:
                 epoch=epoch,
                 best_metric=self.best_metric,
                 config=self.full_config,
+                scheduler=self.scheduler,
+                global_step=self.global_step,
+                train_cfg=dict(self.cfg),
             )
             score = metrics.get(self.monitor)
             if score is not None and score > self.best_metric:
@@ -193,6 +278,9 @@ class Trainer:
                     epoch=epoch,
                     best_metric=self.best_metric,
                     config=self.full_config,
+                    scheduler=self.scheduler,
+                    global_step=self.global_step,
+                    train_cfg=dict(self.cfg),
                 )
                 logger.info("epoch %d 刷新最佳 %s=%.4f，已保存 best.pth", epoch, self.monitor, score)
 

tests/test_train_smoke.py

@@ -13,6 +13,7 @@ pytest.importorskip("yaml")
 pytest.importorskip("PIL")
 
 import numpy as np  # noqa: E402
+import torch  # noqa: E402
 import yaml  # noqa: E402
 from PIL import Image  # noqa: E402
 
@@ -150,3 +151,67 @@ def test_device_flag_overrides_config(project, monkeypatch, tmp_path):
 
 def config_path_of(project):
     return project[0]
+
+
+def test_resume_takes_hyperparams_from_config(project, monkeypatch, tmp_path):
+    """检测训练续训：lr / weight_decay / 总轮数以 yaml 为准，检查点只提供进度。
+
+    回归测试：load_checkpoint 直接 optimizer.load_state_dict，把超参连同动量一起还原，
+    改了配置再续训会被静默覆盖；而且此前根本没保存调度器状态，学习率调度会从头开始。
+    """
+    import math
+
+    import yaml as yaml_mod
+
+    config_path, root = project
+    cfg = yaml_mod.safe_load(config_path.read_text(encoding="utf-8"))
+    cfg["train"].update({"epochs": 4, "optimizer": {"name": "adamw", "lr": 1e-3, "weight_decay": 1e-4},
+                         "scheduler": {"name": "cosine"}})
+    config_path.write_text(yaml_mod.safe_dump(cfg, allow_unicode=True), encoding="utf-8")
+
+    work_dir = tmp_path / "runs"
+    assert run(train_mod, ["train.py", "--config", str(config_path), "--work-dir", str(work_dir)],
+               monkeypatch) == 0
+
+    captured = {}
+    real_fit = train_mod.Trainer.fit
+
+    def probe(self):
+        captured["groups"] = [dict(g) for g in self.optimizer.param_groups]
+        captured["scheduler"] = self.scheduler
+        captured["start_epoch"] = self.start_epoch
+        return 0.0
+
+    monkeypatch.setattr(train_mod.Trainer, "fit", probe)
+    assert run(
+        train_mod,
+        ["train.py", "--config", str(config_path), "--work-dir", str(work_dir),
+         "--resume", str(work_dir / "last.pth"), "--set",
+         "train.epochs=10", "train.optimizer.lr=5e-4", "train.optimizer.weight_decay=0.05"],
+        monkeypatch,
+    ) == 0
+    monkeypatch.setattr(train_mod.Trainer, "fit", real_fit)
+
+    groups = captured["groups"]
+    assert groups[0]["initial_lr"] == pytest.approx(5e-4)     # 不是检查点里的 1e-3
+    assert groups[0]["weight_decay"] == pytest.approx(0.05)   # 不是检查点里的 1e-4
+    assert groups[1]["weight_decay"] == 0.0                   # BN/bias 分组结构没被冲掉
+
+    scheduler = captured["scheduler"]
+    assert scheduler.T_max == 10
+    assert scheduler.last_epoch == captured["start_epoch"] == 4  # 调度进度已快进，不是从 0 开始
+    assert groups[0]["lr"] == pytest.approx(5e-4 * (1 + math.cos(math.pi * 4 / 10)) / 2)
+
+
+def test_checkpoint_carries_scheduler_and_progress(project, monkeypatch, tmp_path):
+    """检查点要带上调度器状态与 global_step，否则续训无从恢复调度进度。"""
+    config_path, _ = project
+    work_dir = tmp_path / "runs"
+    assert run(train_mod, ["train.py", "--config", str(config_path), "--work-dir", str(work_dir)],
+               monkeypatch) == 0
+
+    state = torch.load(work_dir / "last.pth", map_location="cpu", weights_only=False)
+    expected = {"model", "optimizer", "scheduler", "epoch", "global_step", "best_metric", "train_cfg"}
+    assert expected <= set(state)
+    assert state["global_step"] > 0
+    assert state["train_cfg"]["optimizer"]["lr"] is not None

train.py

@@ -87,12 +87,7 @@ def main() -> int:
     )
 
     if args.resume:
-        from engine import load_checkpoint
+        trainer.resume(args.resume)
-
-        state = load_checkpoint(args.resume, model, trainer.optimizer, map_location=device)
-        trainer.start_epoch = int(state.get("epoch", 0)) + 1
-        trainer.best_metric = float(state.get("best_metric") or float("-inf"))
-        logger.info("已从 %s 恢复，起始 epoch=%d", args.resume, trainer.start_epoch)
 
     cfg.dump(work_dir / "config.yaml")
     best = trainer.fit()