fix(simclr): 续训时改变总轮数不再让学习率卡在 0

resume 会恢复调度器状态，其中包含存档时的 T_max。原计划 10 轮的 cosine 跑完后
学习率已降到 0，此时把 epochs 改成 30 接着练，学习率仍是 0——训练照常进行、
日志照常打印，但一步都学不动。是静默失败，比报错更难发现。

- 总轮数与存档一致时行为不变，直接恢复调度器状态
- 不一致时按新的总轮数重建调度器，并快进到当前 epoch；重建前先把 lr 复位回
  initial_lr，因为 cosine 的链式实现基于当前 lr 递推，从 0 出发只会一路留在 0
- 实测续训到 30 轮时 lr 恢复为 7.50e-04，与 cosine 在 10/30 处的理论值一致
- 补三个回归测试：改变总轮数后的学习率、总轮数不变时直接恢复状态、
  save_interval 落的中间盘可直接续训
- 修正一个旧测试：它断言续训后 lr 仍等于旧计划末尾的 0，编码的是错误行为

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

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)

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"