feat(simclr): 支持断点续训，恢复优化器与调度器状态

此前 simclr.checkpoint 的语义是"有权重就跳过训练"，没有真正的续训：中途停掉再启动只能从头再来，或者放弃已训的轮次。 - 新增 simclr_last.pth：骨架 + 投影头 + 优化器 + 调度器 + scaler + epoch + step，与给下游用的 simclr_backbone.pth 并存，两者在同一时机落盘 - 新增 simclr.resume：载入完整状态从下一个 epoch 接着练。只恢复权重不恢复动量与 cosine 进度的话，续训时损失会先反弹一段 - resume 与 checkpoint 互斥（同时给直接报错）；拿只含骨架的 simclr_backbone.pth 去 resume 会给出可操作的报错，提示改用 checkpoint - 续训轮数已达计划值时不再训练，只打警告 Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
2026-09-03 09:34:55 +08:00
parent be6173c746
commit 345dee50da
6 changed files with 254 additions and 13 deletions
@@ -194,6 +194,125 @@ def test_simclr_saves_intermediate_checkpoints(tiny_image_dir, tmp_path):
    assert torch.load(path, map_location="cpu", weights_only=False)["epochs"] == 4


+def _make_trainer(tiny_image_dir, work_dir, epochs=4, **cfg_extra):
+    from torch.utils.data import DataLoader
+
+    from datasets.unlabeled import UnlabeledImageDataset, unlabeled_collate
+
+    dataset = UnlabeledImageDataset(
+        tiny_image_dir,
+        transforms=[{"name": "Resize", "size": [32, 32]}, {"name": "ToTensor"}],
+        views=2,
+    )
+    loader = DataLoader(dataset, batch_size=4, shuffle=False, drop_last=True, collate_fn=unlabeled_collate)
+    torch.manual_seed(0)
+    backbone = BACKBONES.build({"name": "tiny_conv", "widths": [8, 16], "depths": [1, 1]})
+    cfg = {
+        "epochs": epochs,
+        "amp": False,
+        "log_interval": 100,
+        "save_interval": 1,
+        "projection": {"hidden_channels": 16, "out_channels": 8},
+        "scheduler": {"name": "cosine"},
+        "swanlab": {"enabled": False},
+        **cfg_extra,
+    }
+    return SimCLRTrainer(backbone, loader, cfg, torch.device("cpu"), work_dir,
+                         resume=cfg_extra.pop("resume", None))
+
+
+def test_save_state_contains_full_training_state(tiny_image_dir, tmp_path):
+    trainer = _make_trainer(tiny_image_dir, tmp_path / "run", epochs=1)
+    trainer.fit()
+
+    state_path = tmp_path / "run" / "simclr_last.pth"
+    assert state_path.exists()
+    state = torch.load(state_path, map_location="cpu", weights_only=False)
+    # 续训需要的东西一个都不能少，否则动量和学习率进度会被清零
+    assert {"backbone", "projector", "optimizer", "scheduler", "scaler", "epoch", "global_step"} <= set(state)
+    assert state["epoch"] == 0
+
+
+def test_resume_continues_from_the_saved_epoch(tiny_image_dir, tmp_path):
+    """resume 是接着练，不是从头再来。"""
+    work_dir = tmp_path / "run"
+    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]}),
+        first.dataloader,
+        {"epochs": 5, "amp": False, "log_interval": 100, "save_interval": 1,
+         "projection": {"hidden_channels": 16, "out_channels": 8},
+         "scheduler": {"name": "cosine"}, "swanlab": {"enabled": False}},
+        torch.device("cpu"),
+        work_dir,
+        resume=work_dir / "simclr_last.pth",
+    )
+    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)  # 学习率进度接上
+
+    resumed.fit()
+    state = torch.load(work_dir / "simclr_last.pth", map_location="cpu", weights_only=False)
+    assert state["epoch"] == 4                         # 一路练到第 5 轮
+    assert resumed.global_step > steps_after_2
+
+
+def test_resume_rejects_backbone_only_file(tiny_image_dir, tmp_path):
+    """simclr_backbone.pth 只有骨架，不能当续训状态用，要给出可操作的报错。"""
+    work_dir = tmp_path / "run"
+    _make_trainer(tiny_image_dir, work_dir, epochs=1).fit()
+
+    with pytest.raises(ValueError, match="只含骨架权重"):
+        SimCLRTrainer(
+            BACKBONES.build({"name": "tiny_conv", "widths": [8, 16], "depths": [1, 1]}),
+            _make_trainer(tiny_image_dir, work_dir, epochs=1).dataloader,
+            {"epochs": 2, "amp": False, "projection": {"hidden_channels": 16, "out_channels": 8},
+             "swanlab": {"enabled": False}},
+            torch.device("cpu"),
+            work_dir,
+            resume=work_dir / "simclr_backbone.pth",
+        )
+
+    with pytest.raises(FileNotFoundError, match="不存在"):
+        _make_trainer(tiny_image_dir, work_dir, epochs=1, resume=work_dir / "nope.pth")
+
+
+def test_resume_beyond_planned_epochs_does_not_train(tiny_image_dir, tmp_path):
+    work_dir = tmp_path / "run"
+    _make_trainer(tiny_image_dir, work_dir, epochs=3).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": 3, "amp": False, "projection": {"hidden_channels": 16, "out_channels": 8},
+         "swanlab": {"enabled": False}},
+        torch.device("cpu"),
+        work_dir,
+        resume=work_dir / "simclr_last.pth",
+    )
+    steps_before = resumed.global_step
+    resumed.fit()
+    assert resumed.global_step == steps_before  # 一步都没再训
+
+
+def test_resume_and_checkpoint_are_mutually_exclusive(tiny_image_dir, tmp_path, monkeypatch):
+    sys.path.insert(0, "scripts")
+    import cluster_unlabeled as cu
+
+    with pytest.raises(ValueError, match="只能二选一"):
+        cu.build_encoder(
+            {"method": "simclr",
+             "simclr": {"resume": "a.pth", "checkpoint": "b.pth"},
+             "dataset": {"name": "unlabeled", "root": str(tiny_image_dir)}},
+            torch.device("cpu"),
+            tmp_path,
+        )
+
+
 def test_cluster_script_simclr_branch_end_to_end(tiny_image_dir, tmp_path, monkeypatch):
    """method=simclr 全流程：训练 -> 存权重 -> 抽特征 -> 聚类 -> 产出。"""
    sys.path.insert(0, "scripts")