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>

README.md

@@ -223,7 +223,9 @@ uv run python scripts/cluster_unlabeled.py --config configs/cluster/film.yaml \
 
 ```
 runs/simclr/
-├── simclr_backbone.pth      # 每 save_interval 个 epoch 覆盖一次；Ctrl-C 也会先落盘
+├── simclr_backbone.pth      # 只含骨架，给下游用（聚类 / 初始化 DefectNet）
+├── simclr_last.pth          # 完整训练状态，给断点续训用
+├── swanlog/                 # SwanLab 日志
 ├── cluster.log
 ├── assignments.csv
 ├── summary.json
@@ -241,6 +243,26 @@ nohup uv run python scripts/cluster_unlabeled.py --config configs/cluster/film.y
 tail -f runs/simclr.log
 ```
 
+#### 断点续训
+
+训练中断（Ctrl-C、kill、机器重启）后接着练，用 `simclr.resume` 指向 `simclr_last.pth`：
+
+```bash
+uv run python scripts/cluster_unlabeled.py --config configs/cluster/film.yaml \
+    --work-dir runs/simclr \
+    --set method=simclr simclr.epochs=40 loader.num_workers=8 \
+          simclr.resume=runs/simclr/simclr_last.pth
+```
+
+日志会打印"已载入训练状态 …：从 epoch N 接着练到 40"。恢复的不只是权重，还有**投影头、
+优化器动量、学习率调度进度和 step 计数**——少了这些，续训等于把动量和 cosine 进度清零，
+损失会先反弹一段。
+
+> `resume` 与 `checkpoint` 语义不同，二选一（同时给会报错）：
+> **`resume`** = 载入完整状态**接着训练**（要 `simclr_last.pth`）；
+> **`checkpoint`** = 有权重就**跳过训练**直接聚类（要 `simclr_backbone.pth`）。
+> 拿 `simclr_backbone.pth` 去 `resume` 会报错，因为它不含投影头和优化器状态。
+
 #### 第二步：复用权重重新聚类
 
 想换簇数或改聚类参数时**不必重训**——把权重路径填进 `simclr.checkpoint` 就会跳过训练
@@ -285,7 +307,9 @@ simclr:
   temperature: 0.2          # 越小越关注难负样本
   amp: auto                 # CUDA 上开、MPS 上关
   save_interval: 5          # 每 5 个 epoch 落一次盘
-  checkpoint: null          # 填 runs/simclr/simclr_backbone.pth 即可跳过重复训练
+  # 下面两个二选一，语义不同：
+  checkpoint: null          # 有权重就跳过训练直接聚类，例 runs/simclr/simclr_backbone.pth
+  resume: null              # 载入完整状态接着训练，例 runs/simclr/simclr_last.pth
   projection: {hidden_channels: 256, out_channels: 128}
   optimizer: {name: adamw, lr: 0.001, weight_decay: 0.0001}
   scheduler: {name: cosine}

configs/cluster/film.yaml

@@ -110,7 +110,11 @@ simclr:
     tags: [simclr, tiny_conv]
 
   log_interval: 20
-  checkpoint: null              # 填已有的 simclr_backbone.pth 可跳过重复训练
+  # checkpoint 与 resume 二选一，语义不同：
+  #   checkpoint: 有权重就**跳过训练**，直接拿它抽特征聚类（换簇数时用这个）
+  #   resume:     载入骨架+投影头+优化器+调度器，从上次的 epoch **接着训练**
+  checkpoint: null              # 例：runs/simclr/simclr_backbone.pth
+  resume: null                  # 例：runs/simclr/simclr_last.pth
   projection: {hidden_channels: 256, out_channels: 128}
   optimizer: {name: adamw, lr: 0.001, weight_decay: 0.0001}
   scheduler: {name: cosine}

docs/unsupervised-clustering.md

@@ -740,6 +740,7 @@ $k=4$ 稳定性略低于 ImageNet（0.83 vs 0.98），符合"数据的主结构
 | 端到端聚类 | `engine/cluster.py::cluster_features` |
 | NT-Xent 损失 | `engine/simclr.py::nt_xent_loss` |
 | SimCLR 训练循环 | `engine/simclr.py::SimCLRTrainer` |
+| 断点续训（完整状态存取） | `engine/simclr.py::save_state` / `load_state` |
 | 投影头 $g_\varphi$ | `models/heads/projection.py::ProjectionHead` |
 | 设备解析（cuda/mps/cpu，auto 优先级） | `utils/device.py::resolve_device` |
 | AMP 随后端的默认值 | `utils/device.py::resolve_amp` / `default_amp` |

engine/simclr.py

@@ -49,6 +49,9 @@ class SimCLRTrainer:
     cfg: `epochs / temperature / optimizer / scheduler / amp / log_interval / save_interval`
         以及 `swanlab`（可视化，字段同 `train.swanlab`）。
     full_config: 记录到 SwanLab 的超参，一般直接传完整 yaml 配置。
+    resume: 训练状态文件（`simclr_last.pth`）路径，从上次中断处接着练。
+        与 `simclr.checkpoint` 不同：checkpoint 是"有权重就跳过训练直接聚类"，
+        resume 是"载入骨架 + 投影头 + 优化器 + 调度器 + 起始 epoch 继续训练"。
     """
 
     def __init__(
@@ -59,6 +62,7 @@ class SimCLRTrainer:
         device: torch.device,
         work_dir: str | Path,
         full_config: dict | None = None,
+        resume: str | Path | None = None,
     ) -> None:
         self.backbone = backbone
         self.dataloader = dataloader
@@ -86,6 +90,10 @@ class SimCLRTrainer:
         self.peak_memory = PeakMemory(device)
 
         self.global_step = 0
+        self.start_epoch = 0
+        if resume:
+            self.load_state(resume)
+
         self.visualizer = SwanLabLogger(
             work_dir=self.work_dir,
             config=full_config,
@@ -156,35 +164,111 @@ class SimCLRTrainer:
         return meter.avg
 
     def save_backbone(self, epoch: int) -> Path:
-        """只存骨架权重，不含投影头——投影头训练成对增强不变，会丢掉下游有用的信息。"""
+        """只存骨架权重，不含投影头——投影头训练成对增强不变，会丢掉下游有用的信息。
+
+        这是给下游用的产物（聚类、初始化 DefectNet）；要接着训练请用 `simclr_last.pth`。
+        """
         path = self.work_dir / "simclr_backbone.pth"
         torch.save({"model": self.backbone.state_dict(), "epochs": epoch + 1}, path)
         return path
 
-    def fit(self) -> Path:
+    def save_state(self, epoch: int) -> Path:
-        """训练并保存骨架权重（不含投影头），返回权重路径。
+        """完整训练状态，供 `resume` 断点续训。
 
-        每 `save_interval` 个 epoch 落一次盘，中途被 Ctrl-C 或 kill 掉也能拿到当前权重继续用。
+        比 `simclr_backbone.pth` 多了投影头、优化器、调度器与 scaler——少了这些，
+        续训等于把动量和学习率进度清零，损失会先反弹一段。
+        """
+        path = self.work_dir / "simclr_last.pth"
+        torch.save(
+            {
+                "backbone": self.backbone.state_dict(),
+                "projector": self.projector.state_dict(),
+                "optimizer": self.optimizer.state_dict(),
+                "scheduler": self.scheduler.state_dict() if self.scheduler else None,
+                "scaler": self.scaler.state_dict(),
+                "epoch": epoch,
+                "global_step": self.global_step,
+                "epochs_planned": self.epochs,
+            },
+            path,
+        )
+        return path
+
+    def load_state(self, path: str | Path) -> None:
+        """载入 `save_state` 写出的完整状态，从下一个 epoch 接着练。"""
+        path = Path(path)
+        if not path.exists():
+            raise FileNotFoundError(f"续训状态文件不存在: {path}")
+
+        state = torch.load(path, map_location=self.device, weights_only=False)
+        if "projector" not in state:
+            raise ValueError(
+                f"{path} 只含骨架权重，无法续训。续训要用 simclr_last.pth；"
+                "若只是想跳过训练直接聚类，请改用 simclr.checkpoint。"
+            )
+
+        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))
+        logger.info(
+            "已载入训练状态 %s：从 epoch %d 接着练到 %d", path, self.start_epoch, self.epochs
+        )
+        if self.start_epoch >= self.epochs:
+            logger.warning(
+                "状态里已训到 epoch %d，不小于计划的 %d 轮——不会再训练。"
+                "要继续请调大 simclr.epochs。",
+                self.start_epoch,
+                self.epochs,
+            )
+
+    def _checkpoint(self, epoch: int) -> None:
+        """同时落两份：给下游用的骨架，和给续训用的完整状态。"""
+        self.save_backbone(epoch)
+        self.save_state(epoch)
+
+    def fit(self) -> Path:
+        """训练并保存骨架权重（不含投影头），返回骨架路径。
+
+        每 `save_interval` 个 epoch 落一次盘（Ctrl-C 也会先落盘），
+        同时写 `simclr_backbone.pth`（下游用）与 `simclr_last.pth`（续训用）。
         """
         if len(self.dataloader) < 2:
             raise ValueError("SimCLR 需要多个 batch，请调小 batch_size 或增加样本")
 
         path = self.work_dir / "simclr_backbone.pth"
+        if self.start_epoch >= self.epochs:
+            self.save_backbone(self.epochs - 1)
+            self.visualizer.finish()
+            return path
+
+        epoch = self.start_epoch
         try:
-            for epoch in range(self.epochs):
+            for epoch in range(self.start_epoch, self.epochs):
                 self.train_one_epoch(epoch)
                 if self.scheduler is not None:
                     self.scheduler.step()
                 if self.save_interval and (epoch + 1) % self.save_interval == 0:
-                    self.save_backbone(epoch)
+                    self._checkpoint(epoch)
                     logger.info("已保存中间权重（epoch %d）: %s", epoch, path)
         except KeyboardInterrupt:
-            self.save_backbone(epoch)
+            self._checkpoint(epoch)
-            logger.warning("训练被中断，已保存 epoch %d 的权重: %s", epoch, path)
+            logger.warning(
+                "训练被中断，已保存 epoch %d 的权重: %s；续训用 --set simclr.resume=%s",
+                epoch,
+                path,
+                self.work_dir / "simclr_last.pth",
+            )
             self.visualizer.finish()
             return path
 
-        self.save_backbone(self.epochs - 1)
+        self._checkpoint(self.epochs - 1)
         logger.info("自监督骨架已保存: %s", path)
         self.visualizer.finish()
         return path

scripts/cluster_unlabeled.py

@@ -5,6 +5,7 @@
     uv run python scripts/cluster_unlabeled.py --config configs/cluster/film.yaml
 
     # 改簇数 / 换成在本地数据上自监督训练的特征（Mac 上自动走 MPS）
+    # 断点续训：--set method=simclr simclr.resume=runs/simclr/simclr_last.pth
     uv run python scripts/cluster_unlabeled.py --config configs/cluster/film.yaml \
         --set num_clusters=6 method=simclr
 
@@ -76,7 +77,13 @@ def build_encoder(cfg, device, work_dir):
     stages = simclr_cfg.pop("feature_stages", None)
     backbone = BACKBONES.build(backbone_cfg).to(device)
 
+    resume = simclr_cfg.pop("resume", None)
     checkpoint = simclr_cfg.pop("checkpoint", None)
+    if resume and checkpoint:
+        raise ValueError(
+            "simclr.resume 与 simclr.checkpoint 只能二选一："
+            "resume=接着训练，checkpoint=跳过训练直接聚类。"
+        )
     if checkpoint and Path(checkpoint).exists():
         state = torch.load(checkpoint, map_location=device, weights_only=False)
         backbone.load_state_dict(state["model"])
@@ -95,7 +102,9 @@ def build_encoder(cfg, device, work_dir):
         collate_fn=unlabeled_collate,
     )
     full_config = cfg.to_dict() if hasattr(cfg, "to_dict") else dict(cfg)
-    SimCLRTrainer(backbone, loader, simclr_cfg, device, work_dir, full_config=full_config).fit()
+    SimCLRTrainer(
+        backbone, loader, simclr_cfg, device, work_dir, full_config=full_config, resume=resume
+    ).fit()
     return StageFeatureExtractor(backbone, stages).to(device).eval()
 
 

tests/test_simclr.py

@@ -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")