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
+15 -2
View File
@@ -302,8 +302,21 @@ uv run python scripts/cluster_unlabeled.py --config configs/cluster/film.yaml \
```
日志会打印"已载入训练状态 …:从 epoch N 接着练到 40"。恢复的不只是权重,还有**投影头、
优化器动量、学习率调度进度和 step 计数**——少了这些,续训等于把动量和 cosine 进度清零,
损失会先反弹一段。
优化器动量缓冲和 step 计数**——少了这些,续训等于把动量清零,损失会先反弹一段。
> **配置文件是唯一真相源。** 存档只提供"进度"(权重、动量、已训到第几轮);
> `epochs`、`lr`、`weight_decay`、调度器一律以当前 yaml 为准。所以续训时可以顺手调超参:
>
> ```bash
> uv run python scripts/cluster_unlabeled.py --config configs/cluster/film.yaml \
> --work-dir runs/simclr \
> --set method=simclr simclr.epochs=60 simclr.optimizer.lr=5e-4 \
> simclr.resume=runs/simclr/simclr_last.pth
> ```
>
> 与存档不一致时日志会逐项列出(`配置与存档不一致,按当前配置执行:总轮数 40 -> 60;
> optimizer.lr 0.001 -> 0.0005`),不会出现"改了没生效"。学习率调度按新的总轮数重建
> 并快进到当前 epoch,不会沿用存档里的 T_max。
> `resume` 与 `checkpoint` 语义不同,二选一(同时给会报错):
> **`resume`** = 载入完整状态**接着训练**(要 `simclr_last.pth`);
+57 -19
View File
@@ -178,6 +178,8 @@ class SimCLRTrainer:
比 `simclr_backbone.pth` 多了投影头、优化器、调度器与 scaler——少了这些,
续训等于把动量和学习率进度清零,损失会先反弹一段。
同时记下当时的优化器/调度器配置,续训时若与新配置不同会打日志说明按新的来。
"""
path = self.work_dir / "simclr_last.pth"
torch.save(
@@ -190,13 +192,20 @@ class SimCLRTrainer:
"epoch": epoch,
"global_step": self.global_step,
"epochs_planned": self.epochs,
"optimizer_cfg": dict(self.cfg.get("optimizer", {})),
"scheduler_cfg": dict(self.cfg.get("scheduler") or {}),
},
path,
)
return path
def load_state(self, path: str | Path) -> None:
"""载入 `save_state` 写出的完整状态,从下一个 epoch 接着练。"""
"""载入 `save_state` 写出的完整状态,从下一个 epoch 接着练。
**配置文件是唯一真相源**:存档只提供"进度"——模型权重、优化器动量缓冲、
已训到第几个 epoch、global_step;而学习率、weight_decay、总轮数、调度器
一律以当前 yaml 为准。否则改了配置再续训,改动会被存档静默覆盖掉。
"""
path = Path(path)
if not path.exists():
raise FileNotFoundError(f"续训状态文件不存在: {path}")
@@ -210,13 +219,26 @@ class SimCLRTrainer:
self.backbone.load_state_dict(state["backbone"])
self.projector.load_state_dict(state["projector"])
# optimizer.load_state_dict 会整体覆盖 param_groups,把 lr/weight_decay 等超参
# 一并还原成存档时的值。先把配置构造出来的超参抄下来,加载完再按组写回去,
# 这样只留下动量缓冲(真正的"进度"),超参仍以 yaml 为准。
# 按组写回也保住了 build_optimizer 的分组结构(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)
if state.get("scaler"):
self.scaler.load_state_dict(state["scaler"])
self.start_epoch = int(state["epoch"]) + 1
self.global_step = int(state.get("global_step", 0))
self._restore_scheduler(state)
self._log_config_changes(state)
logger.info(
"已载入训练状态 %s:从 epoch %d 接着练到 %d", path, self.start_epoch, self.epochs
)
@@ -229,22 +251,21 @@ class SimCLRTrainer:
)
def _restore_scheduler(self, state: dict) -> None:
"""恢复学习率调度。
"""按**当前配置**重建学习率调度,再快进到已训到的 epoch。
总轮数没变就直接恢复状态;**变了就必须按新的总轮数重建**,只把它快进到当前 epoch。
否则会沿用存档里的 T_max:原计划 10 轮的 cosine 到第 10 轮学习率已经是 0,
这时把 epochs 改成 30 接着练,学习率仍然是 0——训练照跑,但一步都学不动。
不恢复存档里的调度器状态:它带着当时的 T_max 和 base_lrs,一旦 yaml 里改了
总轮数或学习率就会被静默覆盖。原计划 10 轮的 cosine 跑完后 lr 已是 0,
这时把 epochs 改成 30 接着练,沿用旧 T_max 的话 lr 仍是 0——一步都学不动。
cosine / step / multistep 都是 epoch 的确定性函数,快进与恢复状态等价,
所以统一走重建这一条路,没有第二种行为。
"""
if self.scheduler is None or not state.get("scheduler"):
del state # 存档里的 scheduler 状态刻意不用
if self.scheduler is None:
return
planned = int(state.get("epochs_planned", self.epochs))
if planned == self.epochs:
self.scheduler.load_state_dict(state["scheduler"])
return
# 先把 lr 复位回 initial_lr 再重建:cosine 的链式实现是基于当前 lr 递推的,
# 而 optimizer 状态里恢复出来的 lr 可能已经退到 0,直接快进只会一路留在 0。
# cosine 的链式实现基于当前 lr 递推,必须先复位回 initial_lr,
# 否则从 0 出发只会一路留在 0。
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)
@@ -252,13 +273,30 @@ class SimCLRTrainer:
warnings.simplefilter("ignore") # 快进阶段的 "step before optimizer.step()" 属预期
for _ in range(self.start_epoch):
self.scheduler.step()
logger.warning(
"总轮数由 %d 改成了 %d,学习率调度按新的总轮数重建并快进到 epoch %d"
"(不沿用存档里的 T_max),当前 lr=%.2e",
planned,
self.epochs,
self.start_epoch,
def _log_config_changes(self, state: dict) -> None:
"""存档配置与当前配置不一致时说清楚用的是哪一份,避免"改了没生效"的困惑。"""
changes = []
planned = int(state.get("epochs_planned", self.epochs))
if planned != self.epochs:
changes.append(f"总轮数 {planned} -> {self.epochs}")
for key, archived in (state.get("optimizer_cfg") or {}).items():
current = dict(self.cfg.get("optimizer", {})).get(key)
if current is not None and current != archived:
changes.append(f"optimizer.{key} {archived} -> {current}")
archived_scheduler = (state.get("scheduler_cfg") or {}).get("name")
current_scheduler = (dict(self.cfg.get("scheduler") or {})).get("name")
if 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 _checkpoint(self, epoch: int) -> None:
+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()