fix(simclr): 周期性落盘，中途中断不再丢失训练成果

此前 fit() 只在全部 epoch 跑完后才保存骨架，一次跑到 epoch 12 被停掉后
13 分钟的训练全部作废。

- 每 save_interval（默认 5）个 epoch 保存一次骨架
- 捕获 KeyboardInterrupt，中断时先落盘再返回
- 聚类脚本改用 resolve_device，自监督分支不再有 GPU 守卫

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

engine/simclr.py

@@ -1,7 +1,7 @@
 """SimCLR 自监督预训练：在无标注的薄膜图上学一套特征，供聚类使用。
 
 骨架仍然是项目里的纯卷积 backbone（`tiny_conv` 等），预训练完可以直接拿去
-初始化检测模型。属于"训练"，因此调用方必须先过 `utils.device.require_cuda()`。
+初始化检测模型。支持 cuda / mps / cpu，用 `device: auto` 自动挑最快的后端。
 """
 
 from __future__ import annotations
@@ -19,6 +19,7 @@ from torch.utils.data import DataLoader
 from engine.metrics import AverageMeter
 from engine.optim import build_optimizer, build_scheduler
 from models.heads.projection import ProjectionHead
+from utils.device import PeakMemory, resolve_amp, synchronize
 from utils.logging import get_logger
 
 logger = get_logger("engine.simclr")
@@ -44,7 +45,7 @@ class SimCLRTrainer:
     """Args:
     backbone: 输出特征图的卷积骨架（需带 `out_channels`）。
     dataloader: 每个样本产出 2 个视图的 DataLoader（`UnlabeledImageDataset(views=2)`）。
-    cfg: `epochs / temperature / optimizer / scheduler / amp / log_interval` 等。
+    cfg: `epochs / temperature / optimizer / scheduler / amp / log_interval / save_interval` 等。
     """
 
     def __init__(
@@ -65,6 +66,8 @@ class SimCLRTrainer:
         self.epochs = int(self.cfg.get("epochs", 50))
         self.temperature = float(self.cfg.get("temperature", 0.2))
         self.log_interval = int(self.cfg.get("log_interval", 20))
+        # 每隔若干 epoch 落一次盘：自监督动辄跑几十分钟，中途被打断不该把训练成果全丢掉
+        self.save_interval = int(self.cfg.get("save_interval", 5))
 
         self.projector = ProjectionHead(
             backbone.out_channels, **dict(self.cfg.get("projection", {}))
@@ -74,8 +77,9 @@ class SimCLRTrainer:
             self.model, self.cfg.get("optimizer", {"name": "adamw", "lr": 1e-3, "weight_decay": 1e-4})
         )
         self.scheduler = build_scheduler(self.optimizer, self.cfg.get("scheduler"), self.epochs)
-        self.amp = bool(self.cfg.get("amp", True)) and device.type == "cuda"
-        self.scaler = torch.amp.GradScaler("cuda", enabled=self.amp)
+        self.amp = resolve_amp(self.cfg.get("amp", "auto"), device)
+        self.scaler = torch.amp.GradScaler(device.type, enabled=self.amp)
+        self.peak_memory = PeakMemory(device)
 
     def train_one_epoch(self, epoch: int) -> float:
         self.model.train()
@@ -85,7 +89,7 @@ class SimCLRTrainer:
         for step, batch in enumerate(self.dataloader):
             view1, view2 = (view.to(self.device, non_blocking=True) for view in batch["views"][:2])
 
-            with torch.amp.autocast("cuda", enabled=self.amp):
+            with torch.amp.autocast(self.device.type, enabled=self.amp):
                 z1 = self.projector(self.backbone(view1))
                 z2 = self.projector(self.backbone(view2))
                 loss = nt_xent_loss(z1, z2, self.temperature)
@@ -95,6 +99,7 @@ class SimCLRTrainer:
             self.scaler.step(self.optimizer)
             self.scaler.update()
 
+            self.peak_memory.update()
             meter.update(float(loss.detach()), view1.shape[0])
             if step % self.log_interval == 0:
                 logger.info(
@@ -106,20 +111,46 @@ class SimCLRTrainer:
                     self.optimizer.param_groups[0]["lr"],
                 )
 
-        logger.info("simclr epoch %d 完成 loss=%.4f 耗时 %.1fs", epoch, meter.avg, time.time() - start)
+        synchronize(self.device)
+        peak = self.peak_memory.peak_mb
+        logger.info(
+            "simclr epoch %d 完成 loss=%.4f 耗时 %.1fs 显存峰值 %.0f MB%s",
+            epoch,
+            meter.avg,
+            time.time() - start,
+            peak,
+            "" if self.peak_memory.exact else "（MPS 采样下界）",
+        )
         return meter.avg
 
+    def save_backbone(self, epoch: int) -> Path:
+        """只存骨架权重，不含投影头——投影头训练成对增强不变，会丢掉下游有用的信息。"""
+        path = self.work_dir / "simclr_backbone.pth"
+        torch.save({"model": self.backbone.state_dict(), "epochs": epoch + 1}, path)
+        return path
+
     def fit(self) -> Path:
-        """训练并保存骨架权重（不含投影头），返回权重路径。"""
+        """训练并保存骨架权重（不含投影头），返回权重路径。
+
+        每 `save_interval` 个 epoch 落一次盘，中途被 Ctrl-C 或 kill 掉也能拿到当前权重继续用。
+        """
         if len(self.dataloader) < 2:
             raise ValueError("SimCLR 需要多个 batch，请调小 batch_size 或增加样本")
 
-        for epoch in range(self.epochs):
-            self.train_one_epoch(epoch)
-            if self.scheduler is not None:
-                self.scheduler.step()
-
         path = self.work_dir / "simclr_backbone.pth"
-        torch.save({"model": self.backbone.state_dict(), "epochs": self.epochs}, path)
+        try:
+            for epoch in range(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)
+                    logger.info("已保存中间权重（epoch %d）: %s", epoch, path)
+        except KeyboardInterrupt:
+            self.save_backbone(epoch)
+            logger.warning("训练被中断，已保存 epoch %d 的权重: %s", epoch, path)
+            return path
+
+        self.save_backbone(self.epochs - 1)
         logger.info("自监督骨架已保存: %s", path)
         return path