diff --git a/README.md b/README.md index 32bb143..8642466 100644 --- a/README.md +++ b/README.md @@ -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`）； diff --git a/engine/simclr.py b/engine/simclr.py index 29d8983..0905681 100644 --- a/engine/simclr.py +++ b/engine/simclr.py @@ -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: diff --git a/tests/test_simclr.py b/tests/test_simclr.py index 3877e0f..c688998 100644 --- a/tests/test_simclr.py +++ b/tests/test_simclr.py @@ -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()