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>

configs/eval/baseline.yaml

@@ -1,8 +1,8 @@
 _base_:
   - ../train/baseline.yaml
 
-# 评估允许在 CPU 上跑（部署前的精度回归），训练则必须用 GPU。
-device: cuda
+# auto 按 cuda > mps > cpu 自动选；部署前做精度回归可以显式 device=cpu。
+device: auto
 checkpoint: runs/baseline/best.pth
 
 eval:

configs/train/baseline.yaml

@@ -5,11 +5,11 @@ _base_:
 work_dir: runs/baseline
 seed: 42
 deterministic: false
-device: cuda            # 训练只允许 Nvidia GPU，见 utils/device.py
+device: auto            # auto 按 cuda > mps > cpu 挑最快的可用后端；也可写 cuda / mps / cpu
 
 train:
   epochs: 100
-  amp: true
+  amp: auto             # auto: CUDA 上开、MPS 上关（实测 MPS autocast 比 fp32 慢）；也可写 true / false
   clip_grad: 5.0
   log_interval: 20
   monitor: f1           # 用于挑选 best.pth 的验证指标

datasets/build.py

@@ -5,11 +5,11 @@ from __future__ import annotations
 from collections.abc import Mapping
 from typing import Any
 
-import torch
 from torch.utils.data import DataLoader, Dataset
 
 from datasets.film_defect import collate_fn
 from datasets.registry import DATASETS
+from utils.device import pin_memory_supported, resolve_device
 
 
 def build_dataset(cfg: Mapping[str, Any], **kwargs: Any) -> Dataset:
@@ -24,9 +24,9 @@ def build_dataloader(
     pin_memory: bool | None = None,
     drop_last: bool = False,
 ) -> DataLoader:
-    """`pin_memory` 缺省时只在有 CUDA 设备时开启（CPU/MPS 上开启无意义且会告警）。"""
+    """`pin_memory` 缺省时只在 CUDA 上开启（CPU/MPS 上无收益且会告警）。"""
     if pin_memory is None:
-        pin_memory = torch.cuda.is_available()
+        pin_memory = pin_memory_supported(resolve_device("auto"))
     return DataLoader(
         dataset,
         batch_size=batch_size,

engine/trainer.py

@@ -1,7 +1,4 @@
-"""训练循环。
-
-注意：训练入口必须先通过 `utils.device.require_cuda()` 的 GPU 检查。
-"""
+"""训练循环。支持 cuda / mps / cpu 三种后端。"""
 
 from __future__ import annotations
 
@@ -21,6 +18,7 @@ from engine.losses import DetectionLoss
 from engine.metrics import AverageMeter
 from engine.optim import build_optimizer, build_scheduler
 from engine.visualizer import SwanLabLogger
+from utils.device import PeakMemory, resolve_amp
 from utils.labels import MULTILABEL
 from utils.logging import get_logger
 
@@ -66,8 +64,10 @@ class Trainer:
         self.optimizer = build_optimizer(model, self.cfg.get("optimizer", {"name": "adamw", "lr": 1e-3}))
         self.scheduler = build_scheduler(self.optimizer, self.cfg.get("scheduler"), self.epochs)
 
-        self.amp = bool(self.cfg.get("amp", True)) and device.type == "cuda"
-        self.scaler = torch.amp.GradScaler("cuda", enabled=self.amp)
+        # amp: auto 时随后端（CUDA 开、MPS 关——MPS 上 autocast 实测比 fp32 慢）
+        self.amp = resolve_amp(self.cfg.get("amp", "auto"), device)
+        self.scaler = torch.amp.GradScaler(device.type, enabled=self.amp)
+        self.peak_memory = PeakMemory(device)
 
         self.evaluator = (
             Evaluator(
@@ -97,7 +97,7 @@ class Trainer:
         for step, batch in enumerate(self.train_loader):
             batch = _to_device(batch, self.device)
 
-            with torch.amp.autocast("cuda", enabled=self.amp):
+            with torch.amp.autocast(self.device.type, enabled=self.amp):
                 outputs = self.model(batch["images"])
                 losses = self.criterion(outputs, batch)
 
@@ -135,14 +135,18 @@ class Trainer:
         return {key: meter.avg for key, meter in meters.items()}
 
     def _gpu_memory(self) -> dict[str, float]:
-        """本轮训练的显存峰值（GB）；推理端的详细分析见 scripts/profile_inference.py。"""
-        if self.device.type != "cuda":
+        """本轮训练的显存峰值（GB）；推理端的详细分析见 scripts/profile_inference.py。
+
+        CUDA 读内置峰值计数器；MPS 没有峰值 API，用逐 step 采样的最大值（是下界）；
+        CPU 不统计。
+        """
+        if self.device.type == "cpu":
             return {}
         stats = {
-            "gpu_mem_allocated_gb": torch.cuda.max_memory_allocated(self.device) / 1024**3,
-            "gpu_mem_reserved_gb": torch.cuda.max_memory_reserved(self.device) / 1024**3,
+            "gpu_mem_allocated_gb": self.peak_memory.peak_mb / 1024,
+            "gpu_mem_reserved_gb": self.peak_memory.reserved_mb / 1024,
         }
-        torch.cuda.reset_peak_memory_stats(self.device)
+        self.peak_memory.reset()
         return stats
 
     def fit(self) -> float:

eval.py

@@ -4,7 +4,7 @@
 用法:
     python eval.py --config configs/eval/baseline.yaml --checkpoint runs/baseline/best.pth
 
-与训练不同，评估允许在 CPU 上运行（`--set device=cpu`），便于部署前回归。
+设备同样支持 auto / cuda / mps / cpu；部署前回归可以显式 `--device cpu`。
 """
 
 from __future__ import annotations
@@ -18,7 +18,7 @@ from datasets import build_dataloader, build_dataset, load_classes
 from engine import Evaluator
 from models import build_model
 from utils import check_label_mode, load_config, setup_logging
-from utils.device import resolve_device
+from utils.device import DeviceNotAvailableError, describe_device, resolve_device
 
 
 def parse_args() -> argparse.Namespace:
@@ -27,6 +27,7 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument("--checkpoint", default=None, help="覆盖配置中的 checkpoint")
     parser.add_argument("--split", default=None, help="评估哪个 data 划分，默认取配置的 eval.split")
     parser.add_argument("--set", nargs="*", default=[], metavar="KEY=VALUE")
+    parser.add_argument("--device", default=None, help="覆盖配置中的 device：auto / cuda / mps / cpu")
     parser.add_argument("--out", default=None, help="把指标写入 json 文件")
     return parser.parse_args()
 
@@ -45,7 +46,12 @@ def main() -> int:
         logger.error("必须通过 --checkpoint 或配置中的 checkpoint 指定权重")
         return 1
 
-    device = resolve_device(str(cfg.get("device", "cuda")))
+    try:
+        device = resolve_device(args.device or cfg.get("device", "auto"))
+    except DeviceNotAvailableError as exc:
+        logger.error("%s", exc)
+        return 1
+    logger.info("评估设备: %s", describe_device(device))
     classes = load_classes(cfg["classes"])
 
     # label_mode 由 yaml 决定，模型与数据集必须一致，这里统一下发。

scripts/check_env.py

@@ -1,7 +1,7 @@
 #!/usr/bin/env python
-"""环境自检：确认依赖与 Nvidia GPU 是否就绪。
+"""环境自检：确认依赖与可用的计算后端。
 
-python scripts/check_env.py
+    python scripts/check_env.py
 """
 
 from __future__ import annotations
@@ -11,29 +11,44 @@ from pathlib import Path
 
 sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
 
-from utils.device import NoNvidiaGpuError, cuda_devices, require_cuda  # noqa: E402
+from utils.device import (  # noqa: E402
+    available_backends,
+    default_amp,
+    describe_device,
+    memory_stats,
+    resolve_device,
+)
 
 
 def main() -> int:
     print(f"Python: {sys.version.split()[0]}")
     try:
         import torch
-
-        print(f"PyTorch: {torch.__version__} (CUDA build: {torch.version.cuda})")
     except ImportError:
         print("PyTorch: 未安装 —— 先执行 `uv sync`")
         return 1
 
-    devices = cuda_devices()
-    print(f"Nvidia GPU: {devices if devices else '无'}")
+    print(f"PyTorch: {torch.__version__}  (CUDA 构建: {torch.version.cuda})")
+    print(f"可用后端（按优先级）: {available_backends()}")
 
-    try:
-        require_cuda()
-    except NoNvidiaGpuError as exc:
-        print(f"训练不可用: {exc}")
-        print("（评估/导出仍可在 CPU 上进行）")
+    print("\n后端明细:")
+    print(f"  cuda: available={torch.cuda.is_available()}  device_count={torch.cuda.device_count()}")
+    print(
+        f"  mps : available={torch.backends.mps.is_available()}  "
+        f"built={torch.backends.mps.is_built()}"
+    )
+
+    device = resolve_device("auto")
+    print(f"\ndevice=auto 会选中: {describe_device(device)}")
+    print(f"  AMP 默认: {'开' if default_amp(device) else '关'}")
+    stats = memory_stats(device)
+    if stats:
+        print("  " + "  ".join(f"{k}={v:.0f}" for k, v in stats.items()))
+
+    if device.type == "cpu":
+        print("\n⚠️  只有 CPU 可用，训练会非常慢（实测比 MPS 慢约 30 倍），建议换有 GPU 的机器。")
         return 1
-    print("训练环境就绪。")
+    print("\n训练环境就绪。")
     return 0
 
 

scripts/profile_inference.py

@@ -31,6 +31,12 @@ from torch.profiler import ProfilerActivity, profile, record_function  # noqa: E
 from datasets import load_classes  # noqa: E402
 from models import build_model, decode_heatmap  # noqa: E402
 from utils import load_config, setup_logging  # noqa: E402
+from utils.device import (  # noqa: E402
+    DeviceNotAvailableError,
+    PeakMemory,
+    resolve_device,
+    synchronize,
+)
 
 logger = setup_logging()
 
@@ -45,7 +51,7 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument("--image-size", type=int, nargs=2, default=None, help="默认取配置的 image_size")
     parser.add_argument("--iters", type=int, default=50, help="计时轮数")
     parser.add_argument("--warmup", type=int, default=10, help="预热轮数（不计入统计）")
-    parser.add_argument("--device", default=None, help="默认有 GPU 用 cuda，否则 cpu")
+    parser.add_argument("--device", default=None, help="auto / cuda / mps / cpu，默认 auto")
     parser.add_argument("--amp", action="store_true", help="用 float16 autocast 推理")
     parser.add_argument("--postprocess", action="store_true", help="把 heatmap 解码也计入")
     parser.add_argument("--rows", type=int, default=15, help="算子排行显示条数")
@@ -73,7 +79,7 @@ def build(args: argparse.Namespace, device: torch.device) -> torch.nn.Module:
 @torch.no_grad()
 def run_once(model: torch.nn.Module, images: torch.Tensor, args: argparse.Namespace) -> None:
     device_type = images.device.type
-    with torch.autocast(device_type, dtype=torch.float16, enabled=args.amp and device_type == "cuda"):
+    with torch.autocast(device_type, dtype=torch.float16, enabled=args.amp and device_type != "cpu"):
         with record_function("forward"):
             outputs = model(images)
     if args.postprocess:
@@ -89,19 +95,17 @@ def run_once(model: torch.nn.Module, images: torch.Tensor, args: argparse.Namesp
 def measure(model: torch.nn.Module, images: torch.Tensor, args: argparse.Namespace) -> dict:
     """预热后计时，并统计 CUDA 显存峰值。"""
     device = images.device
-    is_cuda = device.type == "cuda"
+    peak = PeakMemory(device)
 
     for _ in range(args.warmup):
         run_once(model, images, args)
-    if is_cuda:
-        torch.cuda.synchronize(device)
-        torch.cuda.reset_peak_memory_stats(device)
+    synchronize(device)
+    peak.reset()
 
     start = time.perf_counter()
     for _ in range(args.iters):
         run_once(model, images, args)
-    if is_cuda:
-        torch.cuda.synchronize(device)
+    synchronize(device)
     elapsed = time.perf_counter() - start
 
     batch = images.shape[0]
@@ -110,11 +114,12 @@ def measure(model: torch.nn.Module, images: torch.Tensor, args: argparse.Namespa
         "latency_ms": elapsed / args.iters * 1000,
         "throughput_fps": args.iters * batch / elapsed,
     }
-    if is_cuda:
+    if device.type != "cpu":
         result.update(
             {
-                "peak_allocated_mb": torch.cuda.max_memory_allocated(device) / 1024**2,
-                "peak_reserved_mb": torch.cuda.max_memory_reserved(device) / 1024**2,
+                "peak_allocated_mb": peak.peak_mb,
+                "peak_reserved_mb": peak.reserved_mb,
+                "peak_is_exact": peak.exact,
                 "weights_mb": sum(p.numel() * p.element_size() for p in model.parameters()) / 1024**2,
             }
         )
@@ -124,13 +129,13 @@ def measure(model: torch.nn.Module, images: torch.Tensor, args: argparse.Namespa
 def profile_once(model: torch.nn.Module, images: torch.Tensor, args: argparse.Namespace) -> None:
     """算子级 profiling：按显存占用与耗时各排一次序。"""
     is_cuda = images.device.type == "cuda"
+    # torch.profiler 没有 MPS 的 activity，Metal 侧的 kernel 抓不到，只能看 CPU 侧下发耗时
     activities = [ProfilerActivity.CPU] + ([ProfilerActivity.CUDA] if is_cuda else [])
 
     with profile(activities=activities, profile_memory=True, record_shapes=True) as prof:
         for _ in range(max(args.iters // 5, 1)):
             run_once(model, images, args)
-    if is_cuda:
-        torch.cuda.synchronize(images.device)
+    synchronize(images.device)
 
     memory_key = "self_cuda_memory_usage" if is_cuda else "self_cpu_memory_usage"
     time_key = "self_cuda_time_total" if is_cuda else "self_cpu_time_total"
@@ -148,9 +153,18 @@ def profile_once(model: torch.nn.Module, images: torch.Tensor, args: argparse.Na
 
 def main() -> int:
     args = parse_args()
-    device = torch.device(args.device or ("cuda" if torch.cuda.is_available() else "cpu"))
-    if device.type != "cuda":
-        logger.warning("当前在 %s 上运行：没有显存数据，只能看耗时与 CPU 内存。", device)
+    try:
+        device = resolve_device(args.device or "auto")
+    except DeviceNotAvailableError as exc:
+        logger.error("%s", exc)
+        return 1
+    if device.type == "cpu":
+        logger.warning("当前在 CPU 上运行：没有显存数据，只能看耗时与 CPU 内存。")
+    elif device.type == "mps":
+        logger.warning(
+            "MPS 上 torch.profiler 抓不到 Metal kernel（没有对应的 ProfilerActivity），"
+            "算子明细只反映 CPU 侧下发耗时；显存用 torch.mps 采样，峰值是下界。"
+        )
 
     model = build(args, device)
     cfg = model.cfg
@@ -158,7 +172,7 @@ def main() -> int:
     channels = 1 if cfg.get("gray") else 3
 
     logger.info(
-        "模型 %s | 参数量 %.2fM | 输出 stride %d | 输入 %dx%dx%d | device %s | amp=%s",
+        "模型 %s | 参数量 %.2fM | 输出 stride %d | 输入 %dx%dx%d | 设备 %s | amp=%s",
         cfg["model"]["name"],
         sum(p.numel() for p in model.parameters()) / 1e6,
         model.out_stride,

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]

train.py

@@ -4,8 +4,10 @@
 用法:
     python train.py --config configs/train/baseline.yaml
     python train.py --config configs/train/baseline.yaml --set train.epochs=10 train.optimizer.lr=5e-4
+    python train.py --config configs/train/baseline.yaml --device mps
 
-硬性约束：没有 Nvidia GPU 时直接退出（见 CLAUDE.MD / utils/device.py）。
+设备由配置的 `device` 决定，默认 `auto`：按 cuda > mps > cpu 挑最快的可用后端。
+显式指定但不可用时直接报错，不会悄悄退回 CPU。
 """
 
 from __future__ import annotations
@@ -14,13 +16,11 @@ import argparse
 import sys
 from pathlib import Path
 
-import torch
-
 from datasets import build_dataloader, build_dataset, load_classes
 from engine import Trainer
 from models import build_model
 from utils import check_label_mode, load_config, seed_everything, setup_logging
-from utils.device import NoNvidiaGpuError, require_cuda
+from utils.device import DeviceNotAvailableError, describe_device, resolve_device
 
 
 def parse_args() -> argparse.Namespace:
@@ -28,6 +28,7 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument("--config", required=True, help="yaml 配置路径")
     parser.add_argument("--set", nargs="*", default=[], metavar="KEY=VALUE", help="覆盖配置项")
     parser.add_argument("--work-dir", default=None, help="覆盖配置中的 work_dir")
+    parser.add_argument("--device", default=None, help="覆盖配置中的 device：auto / cuda / mps / cpu")
     parser.add_argument("--resume", default=None, help="从检查点恢复训练")
     return parser.parse_args()
 
@@ -41,14 +42,13 @@ def main() -> int:
     logger = setup_logging(work_dir / "train.log")
 
     try:
-        gpus = require_cuda()
-    except NoNvidiaGpuError as exc:
+        device = resolve_device(args.device or cfg.get("device", "auto"))
+    except DeviceNotAvailableError as exc:
         logger.error("%s", exc)
         return 1
-    logger.info("检测到 GPU: %s", ", ".join(gpus))
+    logger.info("训练设备: %s", describe_device(device))
 
     seed_everything(int(cfg.get("seed", 42)), bool(cfg.get("deterministic", False)))
-    device = torch.device(cfg.get("device", "cuda"))
 
     classes = load_classes(cfg["classes"])
     # label_mode 由 yaml 决定，模型与数据集必须一致，这里统一下发。

utils/__init__.py

@@ -1,7 +1,15 @@
 """通用工具：配置、注册表、日志、随机种子、设备检查。"""
 
 from utils.config import Config, load_config
-from utils.device import require_cuda, resolve_device
+from utils.device import (
+    DeviceNotAvailableError,
+    PeakMemory,
+    available_backends,
+    describe_device,
+    memory_stats,
+    resolve_amp,
+    resolve_device,
+)
 from utils.labels import LABEL_MODES, MULTICLASS, MULTILABEL, check_label_mode
 from utils.logging import get_logger, setup_logging
 from utils.registry import Registry
@@ -10,8 +18,13 @@ from utils.seed import seed_everything
 __all__ = [
     "Config",
     "load_config",
-    "require_cuda",
     "resolve_device",
+    "resolve_amp",
+    "describe_device",
+    "available_backends",
+    "memory_stats",
+    "PeakMemory",
+    "DeviceNotAvailableError",
     "LABEL_MODES",
     "MULTILABEL",
     "MULTICLASS",

utils/device.py

@@ -1,54 +1,189 @@
-"""设备检查。
+"""设备解析与能力查询：统一处理 cuda / mps / cpu 三种后端。
 
-项目硬性约定（见 CLAUDE.MD）：禁止在没有 Nvidia GPU 的设备上训练。
-所有训练入口必须先调用 `require_cuda()`，检查不通过直接抛错退出。
+设计要点：
+- `resolve_device("auto")` 按 cuda > mps > cpu 的优先级自动挑最快的后端；
+- 显式指定但不可用时直接报错，而不是悄悄退回 CPU——训练慢 30 倍却不吭声是最坏的情况；
+- AMP 的默认值随后端而变：CUDA 上开（有 tensor core），MPS 上关
+  （实测 tiny_conv @128² 每步 fp32 76 ms、fp16 90 ms，autocast 只剩转换开销反而更慢）；
+- 显存统计跨后端统一：CUDA 有内置峰值计数器，MPS 没有，靠采样取最大值。
 """
 
 from __future__ import annotations
 
+import torch
 
-class NoNvidiaGpuError(RuntimeError):
-    """当前设备不具备可用的 Nvidia GPU。"""
+#: 自动选择时的优先级
+BACKEND_PRIORITY = ("cuda", "mps", "cpu")
 
 
-def cuda_devices() -> list[str]:
-    """返回可用的 Nvidia GPU 名称列表；没有 torch 或没有 GPU 时返回空列表。"""
-    try:
-        import torch
-    except ImportError:
-        return []
-    if not torch.cuda.is_available() or torch.version.cuda is None:
-        return []
-    return [torch.cuda.get_device_name(i) for i in range(torch.cuda.device_count())]
+class DeviceNotAvailableError(RuntimeError):
+    """请求的设备在当前机器上不可用。"""
 
 
-def require_cuda() -> list[str]:
-    """训练前的守卫；无 Nvidia GPU 时抛出 `NoNvidiaGpuError`。"""
-    try:
-        import torch
-    except ImportError as exc:
-        raise NoNvidiaGpuError("未安装 PyTorch，无法确认 Nvidia GPU，训练已中止。") from exc
-
-    if torch.version.cuda is None:
-        raise NoNvidiaGpuError(
-            f"当前 PyTorch ({torch.__version__}) 不是 CUDA 构建（可能是 CPU/ROCm/MPS 版本），"
-            "本项目禁止在非 Nvidia GPU 设备上训练。"
-        )
-    if not torch.cuda.is_available() or torch.cuda.device_count() == 0:
-        raise NoNvidiaGpuError(
-            "未检测到可用的 Nvidia GPU，本项目禁止在无 GPU 的设备上训练。"
-            "评估/推理请使用 eval.py 并显式指定 device=cpu。"
-        )
-    return cuda_devices()
+def _cuda_ready() -> bool:
+    return torch.cuda.is_available() and torch.cuda.device_count() > 0
 
 
-def resolve_device(device: str = "cuda"):
-    """把配置中的 device 字符串解析为 `torch.device`。
+def _mps_ready() -> bool:
+    return torch.backends.mps.is_available() and torch.backends.mps.is_built()
 
-    仅用于评估/推理等允许 CPU 的场景；训练路径请先走 `require_cuda()`。
+
+def available_backends() -> list[str]:
+    """按优先级返回当前可用的后端，末尾必定有 cpu。"""
+    backends = []
+    if _cuda_ready():
+        backends.append("cuda")
+    if _mps_ready():
+        backends.append("mps")
+    backends.append("cpu")
+    return backends
+
+
+def resolve_device(spec: str = "auto") -> torch.device:
+    """把配置里的 device 字符串解析成 `torch.device`。
+
+    Args:
+        spec: "auto" 自动挑最优后端；也可写 "cuda" / "cuda:1" / "mps" / "cpu"。
+
+    Raises:
+        DeviceNotAvailableError: 显式请求的后端不可用。
     """
-    import torch
+    spec = str(spec).strip().lower()
+    if spec in ("auto", ""):
+        return torch.device(available_backends()[0])
 
-    if device.startswith("cuda") and not torch.cuda.is_available():
-        raise NoNvidiaGpuError(f"请求了 {device}，但当前设备没有可用的 Nvidia GPU。")
-    return torch.device(device)
+    backend = spec.split(":", 1)[0]
+    if backend == "cuda" and not _cuda_ready():
+        raise DeviceNotAvailableError(
+            f"请求了 {spec}，但当前没有可用的 Nvidia GPU"
+            f"（PyTorch {torch.__version__}，CUDA 构建：{torch.version.cuda}）。"
+            f"可用后端：{available_backends()}，或用 device=auto 自动选择。"
+        )
+    if backend == "mps" and not _mps_ready():
+        raise DeviceNotAvailableError(
+            f"请求了 {spec}，但当前 PyTorch 不支持 MPS"
+            f"（available={torch.backends.mps.is_available()}, built={torch.backends.mps.is_built()}）。"
+            f"可用后端：{available_backends()}。"
+        )
+    if backend not in ("cuda", "mps", "cpu"):
+        raise DeviceNotAvailableError(f"未知设备 {spec}，支持 cuda / mps / cpu / auto。")
+    return torch.device(spec)
+
+
+def describe_device(device: torch.device | str) -> str:
+    """给日志用的人类可读描述。"""
+    device = torch.device(device)
+    if device.type == "cuda":
+        index = device.index or 0
+        return f"{torch.cuda.get_device_name(index)} (cuda:{index})"
+    if device.type == "mps":
+        return "Apple Silicon GPU (Metal / MPS)"
+    return "CPU"
+
+
+def amp_supported(device: torch.device | str) -> bool:
+    """该后端是否支持 autocast 混合精度。"""
+    return torch.device(device).type in ("cuda", "mps")
+
+
+def default_amp(device: torch.device | str) -> bool:
+    """AMP 的默认开关：只有 CUDA 默认开。
+
+    MPS 支持 autocast 且数值正常，但实测比 fp32 更慢，所以默认关闭；
+    需要省显存时可以在配置里显式 `amp: true` 打开。
+    """
+    return torch.device(device).type == "cuda"
+
+
+def resolve_amp(setting: bool | str | None, device: torch.device | str) -> bool:
+    """解析配置里的 `amp` 字段：`auto`/`None` 随后端，`true`/`false` 显式指定。"""
+    device = torch.device(device)
+    if setting is None or (isinstance(setting, str) and setting.lower() == "auto"):
+        return default_amp(device)
+    if not bool(setting):
+        return False
+    return amp_supported(device)
+
+
+def pin_memory_supported(device: torch.device | str) -> bool:
+    """只有 CUDA 的锁页内存有意义；MPS 上开启会告警且无收益。"""
+    return torch.device(device).type == "cuda"
+
+
+def synchronize(device: torch.device | str) -> None:
+    """等待该设备上的异步算子完成——计时前后必须调用，否则测到的是下发时间。"""
+    device = torch.device(device)
+    if device.type == "cuda":
+        torch.cuda.synchronize(device)
+    elif device.type == "mps":
+        torch.mps.synchronize()
+
+
+def memory_stats(device: torch.device | str) -> dict[str, float]:
+    """统一的显存查询，单位 MB。
+
+    CUDA 返回 allocated / reserved / peak；MPS 只有当前分配量与驱动占用
+    （Metal 没有暴露峰值计数器）；CPU 返回空字典。
+    """
+    device = torch.device(device)
+    mb = 1024.0 * 1024.0
+    if device.type == "cuda":
+        return {
+            "allocated_mb": torch.cuda.memory_allocated(device) / mb,
+            "reserved_mb": torch.cuda.memory_reserved(device) / mb,
+            "peak_allocated_mb": torch.cuda.max_memory_allocated(device) / mb,
+            "peak_reserved_mb": torch.cuda.max_memory_reserved(device) / mb,
+        }
+    if device.type == "mps":
+        return {
+            "allocated_mb": torch.mps.current_allocated_memory() / mb,
+            "driver_allocated_mb": torch.mps.driver_allocated_memory() / mb,
+        }
+    return {}
+
+
+class PeakMemory:
+    """跨后端的显存峰值统计。
+
+    CUDA 有内置峰值计数器，直接读；MPS 没有，只能在 `update()` 时采样取最大值——
+    因此 MPS 上的峰值是**下界**（采样点之间的瞬时尖峰会被漏掉）。
+    """
+
+    def __init__(self, device: torch.device | str) -> None:
+        self.device = torch.device(device)
+        self._peak = 0.0
+        self.reset()
+
+    def reset(self) -> None:
+        self._peak = 0.0
+        if self.device.type == "cuda":
+            torch.cuda.reset_peak_memory_stats(self.device)
+        self.update()
+
+    def update(self) -> None:
+        if self.device.type != "mps":
+            return
+        self._peak = max(self._peak, torch.mps.current_allocated_memory() / (1024.0 * 1024.0))
+
+    @property
+    def peak_mb(self) -> float:
+        if self.device.type == "cuda":
+            return torch.cuda.max_memory_allocated(self.device) / (1024.0 * 1024.0)
+        if self.device.type == "mps":
+            self.update()
+            return self._peak
+        return 0.0
+
+    @property
+    def reserved_mb(self) -> float:
+        """CUDA 是缓存分配器保留量；MPS 是 Metal 驱动占用；CPU 为 0。"""
+        if self.device.type == "cuda":
+            return torch.cuda.max_memory_reserved(self.device) / (1024.0 * 1024.0)
+        if self.device.type == "mps":
+            return torch.mps.driver_allocated_memory() / (1024.0 * 1024.0)
+        return 0.0
+
+    @property
+    def exact(self) -> bool:
+        """峰值是否精确（MPS 靠采样，只能给下界）。"""
+        return self.device.type != "mps"