From 345dee50da3b9fe19eee609503b5a86c72c0d8fd Mon Sep 17 00:00:00 2001 From: Shaun Zhou Date: Thu, 3 Sep 2026 09:34:55 +0800 Subject: [PATCH] =?UTF-8?q?feat(simclr):=20=E6=94=AF=E6=8C=81=E6=96=AD?= =?UTF-8?q?=E7=82=B9=E7=BB=AD=E8=AE=AD=EF=BC=8C=E6=81=A2=E5=A4=8D=E4=BC=98?= =?UTF-8?q?=E5=8C=96=E5=99=A8=E4=B8=8E=E8=B0=83=E5=BA=A6=E5=99=A8=E7=8A=B6?= =?UTF-8?q?=E6=80=81?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 此前 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 --- README.md | 28 +++++++- configs/cluster/film.yaml | 6 +- docs/unsupervised-clustering.md | 1 + engine/simclr.py | 102 ++++++++++++++++++++++++--- scripts/cluster_unlabeled.py | 11 ++- tests/test_simclr.py | 119 ++++++++++++++++++++++++++++++++ 6 files changed, 254 insertions(+), 13 deletions(-) diff --git a/README.md b/README.md index 633406d..278c733 100644 --- a/README.md +++ b/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} diff --git a/configs/cluster/film.yaml b/configs/cluster/film.yaml index 95a2849..a9663fc 100644 --- a/configs/cluster/film.yaml +++ b/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} diff --git a/docs/unsupervised-clustering.md b/docs/unsupervised-clustering.md index 29df87d..6a2a855 100644 --- a/docs/unsupervised-clustering.md +++ b/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` | diff --git a/engine/simclr.py b/engine/simclr.py index 4ed6eb1..ea4e8b1 100644 --- a/engine/simclr.py +++ b/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) - logger.warning("训练被中断，已保存 epoch %d 的权重: %s", epoch, path) + self._checkpoint(epoch) + 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 diff --git a/scripts/cluster_unlabeled.py b/scripts/cluster_unlabeled.py index bfb67f7..b6faf5f 100644 --- a/scripts/cluster_unlabeled.py +++ b/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() diff --git a/tests/test_simclr.py b/tests/test_simclr.py index bef6088..d1a20d1 100644 --- a/tests/test_simclr.py +++ b/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")