From b6f83fcad0f766f5adcec50bec9b00e93f8653cb Mon Sep 17 00:00:00 2001 From: Shaun Zhou Date: Wed, 2 Sep 2026 21:43:35 +0800 Subject: [PATCH] =?UTF-8?q?fix(simclr):=20=E5=91=A8=E6=9C=9F=E6=80=A7?= =?UTF-8?q?=E8=90=BD=E7=9B=98=EF=BC=8C=E4=B8=AD=E9=80=94=E4=B8=AD=E6=96=AD?= =?UTF-8?q?=E4=B8=8D=E5=86=8D=E4=B8=A2=E5=A4=B1=E8=AE=AD=E7=BB=83=E6=88=90?= =?UTF-8?q?=E6=9E=9C?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 此前 fit() 只在全部 epoch 跑完后才保存骨架，一次跑到 epoch 12 被停掉后 13 分钟的训练全部作废。 - 每 save_interval（默认 5）个 epoch 保存一次骨架 - 捕获 KeyboardInterrupt，中断时先落盘再返回 - 聚类脚本改用 resolve_device，自监督分支不再有 GPU 守卫 Co-Authored-By: Claude Opus 5 --- configs/cluster/film.yaml | 9 +++--- engine/simclr.py | 57 ++++++++++++++++++++++++++++-------- scripts/cluster_unlabeled.py | 26 ++++++++-------- 3 files changed, 62 insertions(+), 30 deletions(-) diff --git a/configs/cluster/film.yaml b/configs/cluster/film.yaml index 67a46ea..1958ca1 100644 --- a/configs/cluster/film.yaml +++ b/configs/cluster/film.yaml @@ -22,12 +22,12 @@ work_dir: runs/cluster seed: 42 -device: cuda # method=pretrained 时可以 --set device=cpu +device: auto # auto 按 cuda > mps > cpu 挑最快的可用后端；也可写 cuda / mps / cpu -num_clusters: 6 # 想分几类改这里（或 --num-clusters 6） +num_clusters: 4 # 想分几类改这里（或 --num-clusters 6） # pretrained: ImageNet 预训练卷积骨架直接抽特征，不训练，CPU 也能跑，几分钟出结果。 -# simclr: 在本目录数据上自监督训练项目自己的卷积骨架，域适配更好，需要 Nvidia GPU。 +# simclr: 在本目录数据上自监督训练项目自己的卷积骨架，域适配更好；Mac 上走 MPS 即可。 method: pretrained # method=pretrained 用的骨架（method=simclr 的骨架配在 simclr.backbone，两者字段完全不同， @@ -97,7 +97,8 @@ simclr: epochs: 100 batch_size: 256 temperature: 0.2 # 越小越关注难负样本 - amp: true + amp: auto # CUDA 上开、MPS 上关 + save_interval: 5 # 每 5 个 epoch 落一次盘，中途中断也能拿到权重 log_interval: 20 checkpoint: null # 填已有的 simclr_backbone.pth 可跳过重复训练 projection: {hidden_channels: 256, out_channels: 128} diff --git a/engine/simclr.py b/engine/simclr.py index 1e7075a..01c627d 100644 --- a/engine/simclr.py +++ b/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 diff --git a/scripts/cluster_unlabeled.py b/scripts/cluster_unlabeled.py index 54b4b2c..0851da8 100644 --- a/scripts/cluster_unlabeled.py +++ b/scripts/cluster_unlabeled.py @@ -4,7 +4,7 @@ # 默认 4 类，ImageNet 预训练特征（不训练，CPU 可跑） uv run python scripts/cluster_unlabeled.py --config configs/cluster/film.yaml - # 改簇数 / 换成在本地数据上自监督训练的特征（需要 Nvidia GPU） + # 改簇数 / 换成在本地数据上自监督训练的特征（Mac 上自动走 MPS） uv run python scripts/cluster_unlabeled.py --config configs/cluster/film.yaml \ --set num_clusters=6 method=simclr @@ -40,7 +40,11 @@ from engine.cluster import ( # noqa: E402 ) from models.registry import BACKBONES # noqa: E402 from utils import load_config, seed_everything, setup_logging # noqa: E402 -from utils.device import NoNvidiaGpuError, require_cuda, resolve_device # noqa: E402 +from utils.device import ( # noqa: E402 + DeviceNotAvailableError, + describe_device, + resolve_device, +) logger = None @@ -49,7 +53,7 @@ def build_encoder(cfg, device, work_dir): """按 `method` 拿一个特征提取器。 pretrained: ImageNet 预训练卷积骨架，纯推理，CPU 也能跑。用顶层的 `backbone`。 - simclr: 在本目录数据上自监督训练一个项目自己的卷积骨架——属于训练，先过 GPU 守卫。 + simclr: 在本目录数据上自监督训练一个项目自己的卷积骨架。 用 `simclr.backbone`，与 pretrained 的骨架配置各自独立：两者的字段完全不同 （arch/pretrained vs widths/depths），共用一个块会在 `_base_` 合并时串味。 """ @@ -62,7 +66,6 @@ def build_encoder(cfg, device, work_dir): if method != "simclr": raise ValueError(f"未知 method: {method}，可选 pretrained / simclr") - require_cuda() # 自监督预训练是训练，禁止在无 Nvidia GPU 的设备上跑 from torch.utils.data import DataLoader from engine.simclr import SimCLRTrainer @@ -200,6 +203,7 @@ def main() -> int: parser.add_argument("--root", default=None, help="覆盖配置里的图片根目录") parser.add_argument("--num-clusters", type=int, default=None, help="覆盖配置里的簇数") parser.add_argument("--work-dir", default=None) + parser.add_argument("--device", default=None, help="覆盖配置中的 device：auto / cuda / mps / cpu") parser.add_argument("--set", nargs="*", default=[], metavar="KEY=VALUE") parser.add_argument("--save-features", action="store_true", help="额外保存 embedding 供复用") parser.add_argument( @@ -220,26 +224,22 @@ def main() -> int: num_clusters = int(cfg.get("num_clusters", 4)) try: - device = resolve_device(str(cfg.get("device", "cuda"))) - except NoNvidiaGpuError as exc: + device = resolve_device(args.device or cfg.get("device", "auto")) + except DeviceNotAvailableError as exc: logger.error("%s", exc) return 1 dataset = build_dataset(dict(cfg["dataset"])) logger.info( - "扫描到 %d 张图，来自 %d 个目录；簇数 %d，method=%s，device=%s", + "扫描到 %d 张图，来自 %d 个目录；簇数 %d，method=%s，设备 %s", len(dataset), len({dataset.source_dir(path) for path in dataset.images}), num_clusters, cfg.get("method", "pretrained"), - device, + describe_device(device), ) - try: - encoder = build_encoder(cfg, device, work_dir) - except NoNvidiaGpuError as exc: - logger.error("%s", exc) - return 1 + encoder = build_encoder(cfg, device, work_dir) from torch.utils.data import DataLoader