test: 单元测试与端到端冒烟

88 个测试覆盖配置继承、注册表、GPU 守卫、YOLO 格式解析、数据集、模型形状、标签模式、聚类原语与 SimCLR 损失，以及 train.py / eval.py / 聚类脚本的入口接线。含两条回归测试： - PCA 白化放大噪声方向导致簇塌陷（默认 whiten=False 的依据） - classes 写成 data.yaml 路径时 num_classes 必须是类别数而非路径字符串长度 Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
2026-09-02 17:55:47 +08:00
parent f433e5cf59
commit d6adea5b05
15 changed files with 1310 additions and 0 deletions
@@ -0,0 +1,148 @@
+"""SimCLR 自监督：损失的数学性质 + 聚类脚本的 simclr 分支端到端。"""
+
+import sys
+
+import pytest
+
+torch = pytest.importorskip("torch")
+pytest.importorskip("PIL")
+
+import numpy as np  # noqa: E402
+from PIL import Image  # noqa: E402
+
+from engine.simclr import SimCLRTrainer, nt_xent_loss  # noqa: E402
+from models.heads.projection import ProjectionHead  # noqa: E402
+from models.registry import BACKBONES  # noqa: E402
+from utils.device import NoNvidiaGpuError  # noqa: E402
+
+
+def test_nt_xent_is_minimal_when_views_agree():
+    """两个视图完全一致且各样本互相正交时，损失应接近理论下界 log(2B−1) 的反面——
+    即正样本对独占概率质量，损失趋近 0。"""
+    z = torch.eye(4)  # 4 个互相正交的样本
+    aligned = nt_xent_loss(z, z.clone(), temperature=0.05)
+    misaligned = nt_xent_loss(z, z.roll(1, dims=0), temperature=0.05)
+
+    assert aligned < 0.1
+    assert misaligned > aligned
+
+
+def test_nt_xent_upper_bound_at_chance_level():
+    """所有表征相同时无法区分正负样本，损失 = log(2B−1)。"""
+    z = torch.ones(8, 16)
+    loss = float(nt_xent_loss(z, z.clone(), temperature=0.2))
+    assert loss == pytest.approx(float(torch.log(torch.tensor(2 * 8 - 1))), abs=1e-4)
+
+
+def test_nt_xent_is_symmetric_and_temperature_sensitive():
+    torch.manual_seed(0)
+    z1, z2 = torch.randn(6, 8), torch.randn(6, 8)
+    assert float(nt_xent_loss(z1, z2, 0.2)) == pytest.approx(float(nt_xent_loss(z2, z1, 0.2)), abs=1e-5)
+    # 温度越低，随机表征的损失越大（softmax 更尖，惩罚更重）
+    assert nt_xent_loss(z1, z2, 0.05) > nt_xent_loss(z1, z2, 0.5)
+
+
+def test_nt_xent_gradient_pulls_positive_pair_together():
+    torch.manual_seed(0)
+    z1 = torch.randn(4, 8, requires_grad=True)
+    z2 = torch.randn(4, 8, requires_grad=True)
+    before = float(torch.nn.functional.cosine_similarity(z1, z2).mean().detach())
+
+    optimizer = torch.optim.SGD([z1, z2], lr=1.0)
+    for _ in range(20):
+        optimizer.zero_grad()
+        nt_xent_loss(z1, z2, 0.2).backward()
+        optimizer.step()
+
+    after = float(torch.nn.functional.cosine_similarity(z1, z2).mean().detach())
+    assert after > before
+
+
+def test_projection_head_accepts_feature_map_and_vector():
+    head = ProjectionHead(16, hidden_channels=32, out_channels=8)
+    assert head(torch.randn(4, 16, 5, 5)).shape == (4, 8)
+    assert head(torch.randn(4, 16)).shape == (4, 8)
+
+
+@pytest.fixture
+def tiny_image_dir(tmp_path):
+    root = tmp_path / "imgs"
+    root.mkdir()
+    rng = np.random.default_rng(0)
+    for i in range(8):
+        Image.fromarray(rng.integers(0, 255, (40, 60, 3), dtype=np.uint8)).save(root / f"{i}.png")
+    return root
+
+
+def test_simclr_trainer_runs_and_saves_backbone_only(tiny_image_dir, tmp_path):
+    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,
+        view_transforms=[{"name": "RandomResizedCrop", "size": [32, 32], "scale": [0.5, 1.0]}],
+    )
+    loader = DataLoader(dataset, batch_size=4, shuffle=True, drop_last=True, collate_fn=unlabeled_collate)
+    backbone = BACKBONES.build({"name": "tiny_conv", "widths": [8, 16], "depths": [1, 1]})
+
+    trainer = SimCLRTrainer(
+        backbone,
+        loader,
+        {"epochs": 2, "amp": False, "log_interval": 100, "projection": {"hidden_channels": 16,
+                                                                        "out_channels": 8}},
+        torch.device("cpu"),
+        tmp_path / "run",
+    )
+    path = trainer.fit()
+
+    assert path.exists()
+    state = torch.load(path, map_location="cpu", weights_only=False)
+    # 只存骨架，不含投影头
+    assert all("projector" not in key and "mlp" not in key for key in state["model"])
+    backbone.load_state_dict(state["model"])  # 能被干净地加载回去
+
+
+def test_cluster_script_simclr_branch_end_to_end(tiny_image_dir, tmp_path, monkeypatch):
+    """method=simclr 全流程：训练 -> 存权重 -> 抽特征 -> 聚类 -> 产出。"""
+    sys.path.insert(0, "scripts")
+    import cluster_unlabeled as cu
+
+    monkeypatch.setattr(cu, "require_cuda", lambda: ["FAKE-GPU"])
+    work_dir = tmp_path / "run"
+    monkeypatch.setattr(
+        sys, "argv",
+        ["cluster_unlabeled.py", "--config", "configs/cluster/film.yaml",
+         "--work-dir", str(work_dir), "--root", str(tiny_image_dir), "--set",
+         "device=cpu", "method=simclr", "loader.num_workers=0", "loader.batch_size=4",
+         "simclr.epochs=1", "simclr.batch_size=4", "simclr.amp=False",
+         "simclr.backbone.widths=[8,16]", "simclr.backbone.depths=[1,1]",
+         "simclr.projection.hidden_channels=16", "simclr.projection.out_channels=8",
+         "num_clusters=2", "n_init=2"],
+    )
+
+    assert cu.main() == 0
+    assert (work_dir / "simclr_backbone.pth").exists()
+    assert (work_dir / "assignments.csv").exists()
+    assert (work_dir / "summary.json").exists()
+    assert (work_dir / "clusters" / "cluster_00").is_dir()
+
+
+def test_simclr_branch_is_blocked_without_gpu(tiny_image_dir, tmp_path, monkeypatch):
+    """无 Nvidia GPU 时 simclr 分支必须被守卫拦下（pretrained 分支不受影响）。"""
+    sys.path.insert(0, "scripts")
+    import cluster_unlabeled as cu
+
+    def refuse():
+        raise NoNvidiaGpuError("没有 Nvidia GPU")
+
+    monkeypatch.setattr(cu, "require_cuda", refuse)
+    monkeypatch.setattr(
+        sys, "argv",
+        ["cluster_unlabeled.py", "--config", "configs/cluster/film.yaml",
+         "--work-dir", str(tmp_path / "run"), "--root", str(tiny_image_dir),
+         "--set", "device=cpu", "method=simclr", "loader.num_workers=0"],
+    )
+    assert cu.main() == 1