diff --git a/README.md b/README.md index 6e2f8ec..d34d9a6 100644 --- a/README.md +++ b/README.md @@ -109,6 +109,19 @@ uv run python scripts/export_onnx.py --config configs/train/resnet18.yaml \ 产物写在 `work_dir`（默认 `runs//`）：`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` 决定，模型的分类头、分类损失、分类指标都跟着切换 diff --git a/engine/checkpoint.py b/engine/checkpoint.py index 87c9334..22e8400 100644 --- a/engine/checkpoint.py +++ b/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, ) diff --git a/engine/trainer.py b/engine/trainer.py index 38ab369..8767812 100644 --- a/engine/trainer.py +++ b/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) diff --git a/tests/test_train_smoke.py b/tests/test_train_smoke.py index ee8fa9f..12dc778 100644 --- a/tests/test_train_smoke.py +++ b/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 diff --git a/train.py b/train.py index 3eedd74..2d5ea6c 100755 --- a/train.py +++ b/train.py @@ -87,12 +87,7 @@ def main() -> int: ) if args.resume: - from engine import load_checkpoint - - 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) + trainer.resume(args.resume) cfg.dump(work_dir / "config.yaml") best = trainer.fit()