diff --git a/engine/simclr.py b/engine/simclr.py index ea4e8b1..29d8983 100644 --- a/engine/simclr.py +++ b/engine/simclr.py @@ -7,6 +7,7 @@ from __future__ import annotations import time +import warnings from collections.abc import Mapping from pathlib import Path from typing import Any @@ -210,13 +211,12 @@ class SimCLRTrainer: self.backbone.load_state_dict(state["backbone"]) self.projector.load_state_dict(state["projector"]) self.optimizer.load_state_dict(state["optimizer"]) - if self.scheduler is not None and state.get("scheduler"): - self.scheduler.load_state_dict(state["scheduler"]) 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) logger.info( "已载入训练状态 %s：从 epoch %d 接着练到 %d", path, self.start_epoch, self.epochs ) @@ -228,6 +228,39 @@ class SimCLRTrainer: self.epochs, ) + def _restore_scheduler(self, state: dict) -> None: + """恢复学习率调度。 + + 总轮数没变就直接恢复状态；**变了就必须按新的总轮数重建**，只把它快进到当前 epoch。 + 否则会沿用存档里的 T_max：原计划 10 轮的 cosine 到第 10 轮学习率已经是 0， + 这时把 epochs 改成 30 接着练，学习率仍然是 0——训练照跑，但一步都学不动。 + """ + if self.scheduler is None or not state.get("scheduler"): + 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。 + 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) + with warnings.catch_warnings(): + 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, + self.optimizer.param_groups[0]["lr"], + ) + def _checkpoint(self, epoch: int) -> None: """同时落两份：给下游用的骨架，和给续训用的完整状态。""" self.save_backbone(epoch) diff --git a/tests/test_simclr.py b/tests/test_simclr.py index 4f0d93d..3877e0f 100644 --- a/tests/test_simclr.py +++ b/tests/test_simclr.py @@ -239,7 +239,6 @@ def test_resume_continues_from_the_saved_epoch(tiny_image_dir, tmp_path): first = _make_trainer(tiny_image_dir, work_dir, epochs=2) first.fit() steps_after_2 = first.global_step - lr_after_2 = first.optimizer.param_groups[0]["lr"] resumed = SimCLRTrainer( BACKBONES.build({"name": "tiny_conv", "widths": [8, 16], "depths": [1, 1]}), @@ -253,7 +252,8 @@ def test_resume_continues_from_the_saved_epoch(tiny_image_dir, tmp_path): ) assert resumed.start_epoch == 2 # 从第 3 轮开始 assert resumed.global_step == steps_after_2 # step 计数接上 - assert resumed.optimizer.param_groups[0]["lr"] == pytest.approx(lr_after_2) # 学习率进度接上 + # 总轮数从 2 改成了 5，学习率按新计划重算（若沿用旧的 T_max，此处会是 0） + assert resumed.optimizer.param_groups[0]["lr"] > 1e-5 resumed.fit() state = torch.load(work_dir / "simclr_last.pth", map_location="cpu", weights_only=False) @@ -261,6 +261,84 @@ def test_resume_continues_from_the_saved_epoch(tiny_image_dir, tmp_path): assert resumed.global_step > steps_after_2 +def test_resume_with_a_longer_schedule_rebuilds_the_lr_curve(tiny_image_dir, tmp_path): + """续训时调大 epochs，学习率必须按新的总轮数重算。 + + 回归测试：此前直接恢复调度器状态，会沿用存档里的 T_max。原计划 10 轮的 cosine + 跑完后 lr 已经是 0，这时把 epochs 改成 30 接着练，lr 仍是 0——训练照跑但一步都学不动。 + """ + import math + + work_dir = tmp_path / "run" + first = _make_trainer(tiny_image_dir, work_dir, epochs=10) + first.fit() + assert first.optimizer.param_groups[0]["lr"] == pytest.approx(0.0, abs=1e-9) # cosine 已走完 + + resumed = SimCLRTrainer( + BACKBONES.build({"name": "tiny_conv", "widths": [8, 16], "depths": [1, 1]}), + first.dataloader, + {"epochs": 30, "amp": False, "log_interval": 999, "save_interval": 1, + "scheduler": {"name": "cosine"}, "optimizer": {"name": "adamw", "lr": 1e-3}, + "projection": {"hidden_channels": 16, "out_channels": 8}, + "swanlab": {"enabled": False}}, + torch.device("cpu"), + work_dir, + resume=work_dir / "simclr_last.pth", + ) + + assert resumed.scheduler.T_max == 30 + lr = resumed.optimizer.param_groups[0]["lr"] + assert lr > 1e-5, "续训后学习率仍是 0，等于白跑" + assert lr == pytest.approx(1e-3 * (1 + math.cos(math.pi * 10 / 30)) / 2, rel=1e-3) + + +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() + + 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"}, + "projection": {"hidden_channels": 16, "out_channels": 8}, + "swanlab": {"enabled": False}}, + torch.device("cpu"), + work_dir, + resume=work_dir / "simclr_last.pth", + ) + assert resumed.scheduler.T_max == 6 + assert resumed.scheduler.last_epoch == 6 + + +def test_intermediate_checkpoint_is_resumable(tiny_image_dir, tmp_path): + """save_interval 落的中间盘必须能直接拿来续训（不必等训练正常结束）。""" + work_dir = tmp_path / "run" + trainer = _make_trainer(tiny_image_dir, work_dir, epochs=10) + # 只手动跑到第 4 轮就落盘，模拟被 kill 掉 + for epoch in range(4): + trainer.train_one_epoch(epoch) + trainer.scheduler.step() + trainer._checkpoint(3) + + resumed = SimCLRTrainer( + BACKBONES.build({"name": "tiny_conv", "widths": [8, 16], "depths": [1, 1]}), + trainer.dataloader, + {"epochs": 10, "amp": False, "log_interval": 999, "save_interval": 1, + "scheduler": {"name": "cosine"}, + "projection": {"hidden_channels": 16, "out_channels": 8}, + "swanlab": {"enabled": False}}, + torch.device("cpu"), + work_dir, + resume=work_dir / "simclr_last.pth", + ) + assert resumed.start_epoch == 4 + resumed.fit() + state = torch.load(work_dir / "simclr_last.pth", map_location="cpu", weights_only=False) + assert state["epoch"] == 9 + + def test_resume_rejects_backbone_only_file(tiny_image_dir, tmp_path): """simclr_backbone.pth 只有骨架，不能当续训状态用，要给出可操作的报错。""" work_dir = tmp_path / "run"