feat: 支持 cuda / mps / cpu 三后端，Mac 上可用 MPS 训练

移除"训练必须有 Nvidia GPU"的硬性守卫，改为按后端能力分派。

- utils/device: resolve_device("auto") 按 cuda > mps > cpu 挑最快的可用后端；
  显式指定不可用的后端时直接报错，不会悄悄退回 CPU——训练慢几十倍却不吭声是最坏的情况
- AMP 随后端取默认值：CUDA 开、MPS 关。MPS 上 autocast 数值正常但实测更慢
  （tiny_conv @128²：fp32 76 ms/步、fp16 90 ms、bf16 108 ms），需要时 amp: true 可强开
- PeakMemory 统一显存峰值统计：CUDA 读内置计数器，MPS 无峰值 API 改用逐 step 采样
  （是下界，日志与产出里已标注）
- profile_inference 支持 MPS，并提示 torch.profiler 没有 MPS 的 ProfilerActivity，
  算子明细只反映 CPU 侧下发耗时
- train.py / eval.py 新增 --device，配置默认改为 auto
- check_env 改为列出全部可用后端与 auto 的选择结果

实测：tiny_conv 训练在 MPS 上比 CPU 快约 32 倍。

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

tests/test_device.py

@@ -0,0 +1,112 @@
+"""设备解析：cuda / mps / cpu 三后端的选择、AMP 默认值与显存统计。"""
+
+import pytest
+
+torch = pytest.importorskip("torch")
+
+from utils.device import (  # noqa: E402
+    BACKEND_PRIORITY,
+    DeviceNotAvailableError,
+    PeakMemory,
+    amp_supported,
+    available_backends,
+    default_amp,
+    describe_device,
+    memory_stats,
+    pin_memory_supported,
+    resolve_amp,
+    resolve_device,
+    synchronize,
+)
+
+
+def test_available_backends_is_ordered_and_always_has_cpu():
+    backends = available_backends()
+    assert backends[-1] == "cpu"
+    assert backends == [b for b in BACKEND_PRIORITY if b in backends]
+    assert len(set(backends)) == len(backends)
+
+
+def test_auto_picks_the_highest_priority_available_backend():
+    device = resolve_device("auto")
+    assert device.type == available_backends()[0]
+    assert resolve_device("") == device  # 空串等同 auto
+
+
+def test_cpu_is_always_resolvable():
+    assert resolve_device("cpu").type == "cpu"
+
+
+def test_explicit_backend_matches_hardware_or_raises():
+    """显式请求不可用的后端必须报错，不能悄悄退回 CPU——那会让训练慢几十倍还不吭声。"""
+    for backend in ("cuda", "mps"):
+        if backend in available_backends():
+            assert resolve_device(backend).type == backend
+        else:
+            with pytest.raises(DeviceNotAvailableError, match=backend):
+                resolve_device(backend)
+
+
+def test_unknown_device_raises():
+    with pytest.raises(DeviceNotAvailableError, match="未知设备"):
+        resolve_device("tpu")
+
+
+def test_amp_defaults_follow_the_backend():
+    # CUDA 有 tensor core，默认开；MPS 上 autocast 实测更慢，默认关；CPU 不支持
+    assert default_amp("cuda") is True
+    assert default_amp("mps") is False
+    assert default_amp("cpu") is False
+    assert amp_supported("cuda") and amp_supported("mps")
+    assert not amp_supported("cpu")
+
+
+def test_resolve_amp_handles_auto_and_explicit():
+    assert resolve_amp("auto", "cuda") is True
+    assert resolve_amp("auto", "mps") is False
+    assert resolve_amp(None, "mps") is False
+    assert resolve_amp(True, "mps") is True  # 显式打开，MPS 支持 autocast
+    assert resolve_amp(True, "cpu") is False  # CPU 不支持，显式打开也无效
+    assert resolve_amp(False, "cuda") is False
+
+
+def test_pin_memory_only_for_cuda():
+    assert pin_memory_supported("cuda")
+    assert not pin_memory_supported("mps")
+    assert not pin_memory_supported("cpu")
+
+
+def test_describe_device_is_readable():
+    assert describe_device("cpu") == "CPU"
+    assert "MPS" in describe_device("mps")
+
+
+def test_memory_stats_shape_per_backend():
+    assert memory_stats("cpu") == {}
+    stats = memory_stats(resolve_device("auto"))
+    device_type = resolve_device("auto").type
+    if device_type == "cuda":
+        assert {"allocated_mb", "reserved_mb", "peak_allocated_mb"} <= set(stats)
+    elif device_type == "mps":
+        assert {"allocated_mb", "driver_allocated_mb"} <= set(stats)
+
+
+def test_peak_memory_tracks_growth_on_accelerators():
+    device = resolve_device("auto")
+    peak = PeakMemory(device)
+
+    if device.type == "cpu":
+        assert peak.peak_mb == 0.0 and peak.reserved_mb == 0.0
+        return
+
+    # MPS 没有峰值 API，靠采样，所以峰值是下界（exact=False）
+    assert peak.exact == (device.type != "mps")
+    buffer = torch.empty(8 * 1024 * 1024 // 4, device=device)  # 约 8 MB
+    peak.update()
+    synchronize(device)
+    assert peak.peak_mb >= 4.0
+    del buffer
+
+
+def test_synchronize_is_a_noop_on_cpu():
+    synchronize("cpu")  # 不应抛异常

tests/test_device_guard.py

@@ -1,26 +0,0 @@
-"""GPU 守卫：项目禁止在没有 Nvidia GPU 的设备上训练。"""
-
-import pytest
-
-from utils.device import NoNvidiaGpuError, cuda_devices, require_cuda
-
-
-def test_require_cuda_matches_hardware():
-    devices = cuda_devices()
-    if devices:
-        assert require_cuda() == devices
-    else:
-        with pytest.raises(NoNvidiaGpuError):
-            require_cuda()
-
-
-def test_resolve_device_rejects_cuda_without_gpu():
-    torch = pytest.importorskip("torch")
-    from utils.device import resolve_device
-
-    if torch.cuda.is_available():
-        assert resolve_device("cuda").type == "cuda"
-    else:
-        with pytest.raises(NoNvidiaGpuError):
-            resolve_device("cuda")
-    assert resolve_device("cpu").type == "cpu"

tests/test_simclr.py

@@ -13,7 +13,6 @@ 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():
@@ -105,12 +104,46 @@ def test_simclr_trainer_runs_and_saves_backbone_only(tiny_image_dir, tmp_path):
     backbone.load_state_dict(state["model"])  # 能被干净地加载回去
 
 
+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
 
-    monkeypatch.setattr(cu, "require_cuda", lambda: ["FAKE-GPU"])
     work_dir = tmp_path / "run"
     monkeypatch.setattr(
         sys, "argv",
@@ -130,19 +163,25 @@ def test_cluster_script_simclr_branch_end_to_end(tiny_image_dir, tmp_path, monke
     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 分支不受影响）。"""
+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
 
-    def refuse():
-        raise NoNvidiaGpuError("没有 Nvidia GPU")
+    from utils.device import describe_device, resolve_device
 
-    monkeypatch.setattr(cu, "require_cuda", refuse)
+    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(tmp_path / "run"), "--root", str(tiny_image_dir),
-         "--set", "device=cpu", "method=simclr", "loader.num_workers=0"],
+         "--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() == 1
+
+    assert cu.main() == 0, f"在 {describe_device(device)} 上跑失败"
+    assert (work_dir / "simclr_backbone.pth").exists()

tests/test_train_smoke.py

@@ -1,7 +1,7 @@
 """train.py / eval.py 的端到端冒烟：YOLO 数据目录 -> 训练 -> 评估。
 
-覆盖入口脚本的接线（classes 解析、label_mode 下发、GPU 守卫），
-所以在 CPU 上跑，`require_cuda` 用 monkeypatch 放行。
+覆盖入口脚本的接线：classes 解析、label_mode 下发、设备解析。固定用 CPU 跑，
+免得测试受机器上有没有 GPU 影响。
 """
 
 import sys
@@ -19,7 +19,6 @@ from PIL import Image  # noqa: E402
 import eval as eval_mod  # noqa: E402
 import train as train_mod  # noqa: E402
 from datasets import write_data_yaml  # noqa: E402
-from utils.device import NoNvidiaGpuError  # noqa: E402
 
 CLASSES = ["scratch", "bubble"]
 TRANSFORMS = [
@@ -85,7 +84,6 @@ def run(module, argv, monkeypatch):
 def test_train_then_eval(project, monkeypatch, tmp_path):
     config_path, root = project
     work_dir = tmp_path / "runs"
-    monkeypatch.setattr(train_mod, "require_cuda", lambda: ["FAKE-GPU"])
 
     assert run(train_mod, ["train.py", "--config", str(config_path), "--work-dir", str(work_dir)],
                monkeypatch) == 0
@@ -115,7 +113,6 @@ def test_classes_from_data_yaml_sizes_the_heads(project, monkeypatch, tmp_path):
         return real_build(cfg)
 
     monkeypatch.setattr(train_mod, "build_model", spy)
-    monkeypatch.setattr(train_mod, "require_cuda", lambda: ["FAKE-GPU"])
 
     run(train_mod, ["train.py", "--config", str(config_path), "--work-dir", str(work_dir)], monkeypatch)
 
@@ -123,15 +120,33 @@ def test_classes_from_data_yaml_sizes_the_heads(project, monkeypatch, tmp_path):
     assert captured["cfg"]["label_mode"] == "multilabel"
 
 
-def test_training_is_blocked_without_nvidia_gpu(project, monkeypatch, tmp_path):
-    config_path, _ = project
+def test_unavailable_device_fails_loudly(project, monkeypatch, tmp_path):
+    """显式请求机器上没有的后端时必须退出，而不是悄悄退回 CPU。"""
+    from utils.device import available_backends
 
-    def refuse():
-        raise NoNvidiaGpuError("没有 Nvidia GPU")
+    missing = next((b for b in ("cuda", "mps") if b not in available_backends()), None)
+    if missing is None:
+        pytest.skip("这台机器 cuda 与 mps 都可用")
 
-    monkeypatch.setattr(train_mod, "require_cuda", refuse)
     assert run(
         train_mod,
-        ["train.py", "--config", str(config_path), "--work-dir", str(tmp_path / "runs3")],
+        ["train.py", "--config", str(config_path_of(project)), "--work-dir", str(tmp_path / "runs3"),
+         "--device", missing],
         monkeypatch,
     ) == 1
+
+
+def test_device_flag_overrides_config(project, monkeypatch, tmp_path):
+    """--device 优先于配置里的 device。"""
+    config_path, _ = project
+    work_dir = tmp_path / "runs4"
+    assert run(
+        train_mod,
+        ["train.py", "--config", str(config_path), "--work-dir", str(work_dir), "--device", "cpu"],
+        monkeypatch,
+    ) == 0
+    assert (work_dir / "best.pth").exists()
+
+
+def config_path_of(project):
+    return project[0]