"""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 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_simclr_logs_metrics_to_swanlab(tiny_image_dir, tmp_path, monkeypatch): """SwanLab 接线：按 step 记 loss/lr，按 epoch 记平均 loss、耗时与显存峰值，结束时收尾。""" from torch.utils.data import DataLoader import engine.simclr as simclr_mod from datasets.unlabeled import UnlabeledImageDataset, unlabeled_collate records: list[tuple[dict, int | None, str]] = [] finished: list[bool] = [] class StubLogger: def __init__(self, **kwargs): self.kwargs = kwargs def log(self, metrics, step=None, prefix=""): records.append((metrics, step, prefix)) def finish(self): finished.append(True) monkeypatch.setattr(simclr_mod, "SwanLabLogger", StubLogger) dataset = UnlabeledImageDataset( tiny_image_dir, transforms=[{"name": "Resize", "size": [32, 32]}, {"name": "ToTensor"}], views=2, ) 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 = simclr_mod.SimCLRTrainer( backbone, loader, {"epochs": 1, "amp": False, "log_interval": 1, "save_interval": 0, "projection": {"hidden_channels": 16, "out_channels": 8}, "swanlab": {"enabled": True, "mode": "disabled"}}, torch.device("cpu"), tmp_path / "run", full_config={"method": "simclr"}, ) trainer.fit() assert records, "没有任何指标被记录" assert all(prefix == "simclr/" for _, _, prefix in records) step_keys = {k for metrics, _, _ in records for k in metrics} assert {"loss", "loss_avg", "lr"} <= step_keys assert {"epoch", "epoch_loss", "epoch_seconds", "gpu_mem_peak_mb"} <= step_keys assert all(step is not None and step > 0 for _, step, _ in records) assert finished == [True] # fit() 结束时调用了 finish() # 完整 yaml 配置作为超参传给了 SwanLab assert trainer.visualizer.kwargs["config"] == {"method": "simclr"} assert trainer.visualizer.kwargs["mode"] == "disabled" def test_simclr_saves_intermediate_checkpoints(tiny_image_dir, tmp_path): """中途落盘：跑到一半被打断也要能拿到权重（此前只在全部 epoch 跑完后才存）。""" 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=True, drop_last=True, collate_fn=unlabeled_collate) backbone = BACKBONES.build({"name": "tiny_conv", "widths": [8, 16], "depths": [1, 1]}) work_dir = tmp_path / "run" trainer = SimCLRTrainer( backbone, loader, {"epochs": 4, "amp": False, "log_interval": 100, "save_interval": 2, "projection": {"hidden_channels": 16, "out_channels": 8}}, torch.device("cpu"), work_dir, ) # 手动走前两个 epoch，模拟"训练中途" trainer.train_one_epoch(0) trainer.train_one_epoch(1) path = trainer.save_backbone(1) assert path.exists() assert torch.load(path, map_location="cpu", weights_only=False)["epochs"] == 2 trainer.fit() # 完整跑完后 epochs 应更新 assert torch.load(path, map_location="cpu", weights_only=False)["epochs"] == 4 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 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_runs_on_the_resolved_device(tiny_image_dir, tmp_path, monkeypatch): """method=simclr 在本机自动选中的后端上（MPS / CUDA / CPU）都应能跑通。""" sys.path.insert(0, "scripts") import cluster_unlabeled as cu from utils.device import describe_device, resolve_device device = resolve_device("auto") work_dir = tmp_path / "auto_run" monkeypatch.setattr( sys, "argv", ["cluster_unlabeled.py", "--config", "configs/cluster/film.yaml", "--work-dir", str(work_dir), "--root", str(tiny_image_dir), "--set", "method=simclr", "loader.num_workers=0", "loader.batch_size=4", "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", "num_clusters=2", "n_init=2"], ) assert cu.main() == 0, f"在 {describe_device(device)} 上跑失败" assert (work_dir / "simclr_backbone.pth").exists()