fix(simclr): 续训以配置为准,不再被存档里的超参静默覆盖

optimizer.load_state_dict 会整体覆盖 param_groups,把 lr / weight_decay 连同动量缓冲
一起还原成存档时的值。改了 yaml 再续训,改动会被静默吃掉:配置写 lr=5e-4,
实际仍按存档的 1e-3 训,scheduler.base_lrs 也是旧值。

改为「配置文件是唯一真相源,存档只提供进度」:

- 加载优化器状态前先抄下配置构造出的每组超参,加载后按组写回,只留下动量缓冲;
  按组写回也保住了 build_optimizer 的分组结构(BN 与 bias 那组 weight_decay=0)
- 调度器统一按当前配置重建再快进到已训 epoch,不再恢复存档状态(它带着旧的
  T_max 与 base_lrs)。cosine/step/multistep 都是 epoch 的确定性函数,
  快进与恢复状态等价,因此不存在第二种行为
- 存档里额外记下当时的 optimizer/scheduler 配置,续训时逐项列出差异,
  明确说明按当前配置执行
- 补回归测试:改 lr/weight_decay/总轮数后续训,三者都按 yaml 生效,
  分组结构与动量缓冲不受影响

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
This commit is contained in:
Shaun Zhou
2026-09-04 21:46:40 +08:00
co-authored by Claude Opus 5
parent 703ba7076b
commit b7505a743d
3 changed files with 112 additions and 23 deletions
+40 -2
View File
@@ -218,7 +218,7 @@ def _make_trainer(tiny_image_dir, work_dir, epochs=4, **cfg_extra):
**cfg_extra,
}
return SimCLRTrainer(backbone, loader, cfg, torch.device("cpu"), work_dir,
resume=cfg_extra.pop("resume", None))
resume=cfg.pop("resume", None))
def test_save_state_contains_full_training_state(tiny_image_dir, tmp_path):
@@ -292,8 +292,46 @@ def test_resume_with_a_longer_schedule_rebuilds_the_lr_curve(tiny_image_dir, tmp
assert lr == pytest.approx(1e-3 * (1 + math.cos(math.pi * 10 / 30)) / 2, rel=1e-3)
def test_resume_takes_hyperparams_from_config_not_checkpoint(tiny_image_dir, tmp_path):
"""续训时 lr / weight_decay / 总轮数一律以 yaml 为准,存档只提供进度。
回归测试:optimizer.load_state_dict 会整体覆盖 param_groups,把超参连同动量缓冲
一起还原成存档时的值——改了配置再续训,改动被静默吃掉。
"""
import math
work_dir = tmp_path / "run"
first = _make_trainer(
tiny_image_dir, work_dir, epochs=4,
optimizer={"name": "adamw", "lr": 1e-3, "weight_decay": 1e-4},
)
first.fit()
resumed = SimCLRTrainer(
BACKBONES.build({"name": "tiny_conv", "widths": [8, 16], "depths": [1, 1]}),
first.dataloader,
{"epochs": 8, "amp": False, "log_interval": 999, "save_interval": 1,
"scheduler": {"name": "cosine"},
"optimizer": {"name": "adamw", "lr": 5e-4, "weight_decay": 0.05},
"projection": {"hidden_channels": 16, "out_channels": 8},
"swanlab": {"enabled": False}},
torch.device("cpu"),
work_dir,
resume=work_dir / "simclr_last.pth",
)
groups = resumed.optimizer.param_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 分组结构没被冲掉
assert resumed.scheduler.base_lrs == pytest.approx([5e-4, 5e-4])
assert resumed.scheduler.T_max == 8
assert groups[0]["lr"] == pytest.approx(5e-4 * (1 + math.cos(math.pi * 4 / 8)) / 2)
assert resumed.optimizer.state, "动量缓冲应当保留——那才是真正要续的'进度'"
def test_resume_with_same_schedule_restores_state(tiny_image_dir, tmp_path):
"""总轮数不变时应当直接恢复调度器状态,不重建。"""
"""总轮数不变时,重建 + 快进的结果应与原调度进度一致。"""
work_dir = tmp_path / "run"
_make_trainer(tiny_image_dir, work_dir, epochs=6).fit()