diff --git a/README.md b/README.md index 30e6ef3..6e2f8ec 100644 --- a/README.md +++ b/README.md @@ -359,6 +359,28 @@ uv run python scripts/cluster_unlabeled.py --config configs/cluster/film.yaml \ > **`checkpoint`** = 有权重就**跳过训练**直接聚类（要 `simclr_backbone.pth`）。 > 拿 `simclr_backbone.pth` 去 `resume` 会报错，因为它不含投影头和优化器状态。 +#### 训练与聚类分开跑 + +自监督训练和聚类在代码上本来就是解耦的（`engine/simclr.py` 与 `engine/cluster.py` +互不引用，唯一接口是 `simclr_backbone.pth`），只是默认一条命令串起来跑。 +需要分两步时——比如在 GPU 机器上专心训练几十小时，之后在别处反复调 k——用 +`--train-only` 断开： + +```bash +# 第一步：只训练，存完权重就退出，不抽特征也不聚类 +uv run python scripts/cluster_unlabeled.py --config configs/cluster/film_simclr_resnet18.yaml \ + --work-dir runs/simclr_r18 --train-only --set simclr.epochs=40 + +# 第二步：只聚类，想换几次 k 就跑几次，每次约 1 分钟 +uv run python scripts/cluster_unlabeled.py --config configs/cluster/film_simclr_resnet18.yaml \ + --work-dir runs/simclr_r18_k4 \ + --set simclr.checkpoint=runs/simclr_r18/simclr_backbone.pth --num-clusters 4 +``` + +第一步只产出 `simclr_backbone.pth` 与 `simclr_last.pth`，没有 `assignments.csv` 之类。 +`--train-only` 遇到 `method=pretrained`（没东西可训）、`simclr.checkpoint`（会跳过训练） +或 `--sweep-k`（要靠聚类才能扫）都会直接报错，不会跑完才发现什么都没产出。 + #### 第二步：复用权重重新聚类 想换簇数或改聚类参数时**不必重训**——把权重路径填进 `simclr.checkpoint` 就会跳过训练 diff --git a/scripts/cluster_unlabeled.py b/scripts/cluster_unlabeled.py index 3951d94..2cf6f81 100644 --- a/scripts/cluster_unlabeled.py +++ b/scripts/cluster_unlabeled.py @@ -389,6 +389,11 @@ def main() -> int: parser.add_argument( "--sweep-k", default=None, metavar="A-B", help="只扫簇数（如 2-8）打印指标表，不产出分组" ) + parser.add_argument( + "--train-only", + action="store_true", + help="只做 SimCLR 自监督训练，存完权重就退出，不抽特征也不聚类", + ) args = parser.parse_args() cfg = load_config(args.config, args.set) @@ -409,18 +414,46 @@ def main() -> int: logger.error("%s", exc) return 1 + method = str(cfg.get("method", "pretrained")) + if args.train_only: + # 这三种组合下 --train-only 要么无事可做、要么与另一个"提前结束"的开关打架， + # 与其跑完才发现什么都没产出，不如立刻说清楚 + if method != "simclr": + logger.error("--train-only 只对 method=simclr 有意义，当前 method=%s（没有要训练的东西）", method) + return 1 + if cfg.get("simclr", {}).get("checkpoint"): + logger.error( + "--train-only 与 simclr.checkpoint 冲突：checkpoint 会跳过训练，" + "加上 --train-only 就什么都不会做。想接着上次练请改用 simclr.resume。" + ) + return 1 + if args.sweep_k: + logger.error("--train-only 与 --sweep-k 不能同时用：前者不聚类，后者要靠聚类才能扫") + return 1 + dataset = build_dataset(dict(cfg["dataset"])) logger.info( "扫描到 %d 张图，来自 %d 个目录；簇数 %d，method=%s，设备 %s", len(dataset), len({dataset.source_dir(path) for path in dataset.images}), num_clusters, - cfg.get("method", "pretrained"), + method, describe_device(device), ) encoder = build_encoder(cfg, device, work_dir) + if args.train_only: + weights = work_dir / "simclr_backbone.pth" + logger.info("--train-only：训练完成，未做聚类。权重在 %s", weights) + logger.info( + "接着聚类：uv run python scripts/cluster_unlabeled.py --config %s " + "--work-dir <新目录> --set simclr.checkpoint=%s", + cfg.get("config_path", "<配置>"), + weights, + ) + return 0 + from torch.utils.data import DataLoader loader = DataLoader( diff --git a/tests/test_simclr.py b/tests/test_simclr.py index 979ff49..05b5a8a 100644 --- a/tests/test_simclr.py +++ b/tests/test_simclr.py @@ -480,6 +480,99 @@ def test_cluster_script_simclr_branch_end_to_end(tiny_image_dir, tmp_path, monke assert (work_dir / "clusters" / "cluster_00").is_dir() +def _cluster_argv(work_dir, image_dir, *extra): + return [ + "cluster_unlabeled.py", "--config", "configs/cluster/film.yaml", + "--work-dir", str(work_dir), "--root", str(image_dir), "--set", + "device=cpu", "loader.num_workers=0", "loader.batch_size=4", + "num_clusters=2", "n_init=2", *extra, + ] + + +def test_train_only_stops_before_clustering(tiny_image_dir, tmp_path, monkeypatch): + """--train-only 只训练存权重，不抽特征也不聚类。""" + sys.path.insert(0, "scripts") + import cluster_unlabeled as cu + + work_dir = tmp_path / "train_only" + monkeypatch.setattr( + sys, "argv", + [*_cluster_argv(work_dir, tiny_image_dir, + "method=simclr", "simclr.epochs=1", "simclr.batch_size=4", + "simclr.backbone.widths=[8,16]", "simclr.backbone.depths=[1,1]", + "simclr.projection.hidden_channels=16", "simclr.projection.out_channels=8"), + "--train-only"], + ) + assert cu.main() == 0 + + assert (work_dir / "simclr_backbone.pth").exists() + assert (work_dir / "simclr_last.pth").exists() # 续训状态照常落盘 + for produced_by_clustering in ("assignments.csv", "summary.json", "cluster_scatter.png", "clusters"): + assert not (work_dir / produced_by_clustering).exists(), produced_by_clustering + + +def test_train_only_then_cluster_only(tiny_image_dir, tmp_path, monkeypatch): + """训练与聚类可以分两次跑：第一次只训练，第二次用 checkpoint 只聚类。""" + sys.path.insert(0, "scripts") + import cluster_unlabeled as cu + + train_dir = tmp_path / "phase1" + monkeypatch.setattr( + sys, "argv", + [*_cluster_argv(train_dir, tiny_image_dir, + "method=simclr", "simclr.epochs=1", "simclr.batch_size=4", + "simclr.backbone.widths=[8,16]", "simclr.backbone.depths=[1,1]", + "simclr.projection.hidden_channels=16", "simclr.projection.out_channels=8"), + "--train-only"], + ) + assert cu.main() == 0 + + cluster_dir = tmp_path / "phase2" + monkeypatch.setattr( + sys, "argv", + _cluster_argv(cluster_dir, tiny_image_dir, + "method=simclr", + f"simclr.checkpoint={train_dir / 'simclr_backbone.pth'}", + "simclr.backbone.widths=[8,16]", "simclr.backbone.depths=[1,1]"), + ) + assert cu.main() == 0 + + assert (cluster_dir / "assignments.csv").exists() + assert not (cluster_dir / "simclr_backbone.pth").exists(), "第二步不该重新训练" + + +@pytest.mark.parametrize( + "extra,reason", + [ + ([], "method=pretrained 没有要训练的东西"), + (["method=simclr", "simclr.checkpoint=whatever.pth"], "checkpoint 会跳过训练"), + ], +) +def test_train_only_rejects_meaningless_combinations( + tiny_image_dir, tmp_path, monkeypatch, extra, reason +): + """这些组合下 --train-only 什么都不会做，应当立刻报错而不是跑完才发现没产出。""" + sys.path.insert(0, "scripts") + import cluster_unlabeled as cu + + monkeypatch.setattr( + sys, "argv", [*_cluster_argv(tmp_path / "bad", tiny_image_dir, *extra), "--train-only"] + ) + assert cu.main() == 1, reason + + +def test_train_only_conflicts_with_sweep_k(tiny_image_dir, tmp_path, monkeypatch): + sys.path.insert(0, "scripts") + import cluster_unlabeled as cu + + monkeypatch.setattr( + sys, "argv", + [*_cluster_argv(tmp_path / "bad", tiny_image_dir, "method=simclr"), + "--train-only", "--sweep-k", "2-4"], + ) + assert cu.main() == 1 + + def test_scatter_plot_is_written_next_to_previews(tiny_image_dir, tmp_path, monkeypatch): """cluster_scatter.png 要和 preview 拼图落在同一级目录。""" sys.path.insert(0, "scripts")