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>
This commit is contained in:
co-authored by
Claude Opus 5
parent
01cfff9b02
commit
105cb38831
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user