fix(simclr): 换优化器种类时不再加载不兼容的动量缓冲

adamw 换成 sgd 时 load_state_dict 不会报错,SGD 会收下 AdamW 的 exp_avg/exp_avg_sq
却永远用不上——动量实际从零开始,但状态里白占内存,之后存盘还会一路带下去。
改为检测到优化器种类变化就跳过加载并给出日志,模型权重与训练进度不受影响。

README 补一张表,说明续训时哪些配置按 yaml、哪些来自存档、哪些不能改。

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
This commit is contained in:
Shaun Zhou
2026-09-05 11:10:04 +08:00
co-authored by Claude Opus 5
parent b7505a743d
commit 004d7a1ce7
3 changed files with 57 additions and 7 deletions
+27
View File
@@ -330,6 +330,33 @@ def test_resume_takes_hyperparams_from_config_not_checkpoint(tiny_image_dir, tmp
assert resumed.optimizer.state, "动量缓冲应当保留——那才是真正要续的'进度'"
def test_resume_across_optimizer_types_drops_incompatible_state(tiny_image_dir, tmp_path):
"""换优化器种类时不加载动量缓冲:AdamW 的 exp_avg 对 SGD 没有意义,硬塞只是白占内存。"""
work_dir = tmp_path / "run"
_make_trainer(
tiny_image_dir, work_dir, epochs=2,
optimizer={"name": "adamw", "lr": 1e-3, "weight_decay": 1e-4},
).fit()
resumed = SimCLRTrainer(
BACKBONES.build({"name": "tiny_conv", "widths": [8, 16], "depths": [1, 1]}),
_make_trainer(tiny_image_dir, work_dir, epochs=1).dataloader,
{"epochs": 6, "amp": False, "log_interval": 999, "save_interval": 1,
"scheduler": {"name": "cosine"},
"optimizer": {"name": "sgd", "lr": 0.01, "momentum": 0.9},
"projection": {"hidden_channels": 16, "out_channels": 8},
"swanlab": {"enabled": False}},
torch.device("cpu"),
work_dir,
resume=work_dir / "simclr_last.pth",
)
assert isinstance(resumed.optimizer, torch.optim.SGD)
assert not resumed.optimizer.state, "不该带着 AdamW 的缓冲"
assert resumed.start_epoch == 2 # 训练进度仍然接上
assert resumed.optimizer.param_groups[0]["momentum"] == 0.9
def test_resume_with_same_schedule_restores_state(tiny_image_dir, tmp_path):
"""总轮数不变时,重建 + 快进的结果应与原调度进度一致。"""
work_dir = tmp_path / "run"