From 141cb907f2bba8030d84714fa2776e1717470bfe Mon Sep 17 00:00:00 2001 From: Shaun Zhou Date: Wed, 2 Sep 2026 21:43:35 +0800 Subject: [PATCH] =?UTF-8?q?feat:=20=E6=94=AF=E6=8C=81=20cuda=20/=20mps=20/?= =?UTF-8?q?=20cpu=20=E4=B8=89=E5=90=8E=E7=AB=AF=EF=BC=8CMac=20=E4=B8=8A?= =?UTF-8?q?=E5=8F=AF=E7=94=A8=20MPS=20=E8=AE=AD=E7=BB=83?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 移除"训练必须有 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 --- configs/eval/baseline.yaml | 4 +- configs/train/baseline.yaml | 4 +- datasets/build.py | 6 +- engine/trainer.py | 28 +++-- eval.py | 12 +- scripts/check_env.py | 41 ++++--- scripts/profile_inference.py | 48 +++++--- tests/test_device.py | 112 ++++++++++++++++++ tests/test_device_guard.py | 26 ----- tests/test_simclr.py | 59 ++++++++-- tests/test_train_smoke.py | 37 ++++-- train.py | 16 +-- utils/__init__.py | 17 ++- utils/device.py | 213 ++++++++++++++++++++++++++++------- 14 files changed, 475 insertions(+), 148 deletions(-) create mode 100644 tests/test_device.py delete mode 100644 tests/test_device_guard.py diff --git a/configs/eval/baseline.yaml b/configs/eval/baseline.yaml index a1d385a..fe22d55 100644 --- a/configs/eval/baseline.yaml +++ b/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: diff --git a/configs/train/baseline.yaml b/configs/train/baseline.yaml index a45bc4b..b06557b 100644 --- a/configs/train/baseline.yaml +++ b/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 的验证指标 diff --git a/datasets/build.py b/datasets/build.py index 298a0a3..e950ba4 100644 --- a/datasets/build.py +++ b/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, diff --git a/engine/trainer.py b/engine/trainer.py index 22e9e1a..38ab369 100644 --- a/engine/trainer.py +++ b/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: diff --git a/eval.py b/eval.py index e88b84e..57e7c2d 100755 --- a/eval.py +++ b/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 决定，模型与数据集必须一致，这里统一下发。 diff --git a/scripts/check_env.py b/scripts/check_env.py index 6a2346d..17a0821 100755 --- a/scripts/check_env.py +++ b/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 diff --git a/scripts/profile_inference.py b/scripts/profile_inference.py index cbf68dc..4d59e9a 100755 --- a/scripts/profile_inference.py +++ b/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, diff --git a/tests/test_device.py b/tests/test_device.py new file mode 100644 index 0000000..4523c33 --- /dev/null +++ b/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") # 不应抛异常 diff --git a/tests/test_device_guard.py b/tests/test_device_guard.py deleted file mode 100644 index 3a55994..0000000 --- a/tests/test_device_guard.py +++ /dev/null @@ -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" diff --git a/tests/test_simclr.py b/tests/test_simclr.py index 40f0223..9a32ccb 100644 --- a/tests/test_simclr.py +++ b/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() diff --git a/tests/test_train_smoke.py b/tests/test_train_smoke.py index e3a7509..ee8fa9f 100644 --- a/tests/test_train_smoke.py +++ b/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] diff --git a/train.py b/train.py index 4575945..3eedd74 100755 --- a/train.py +++ b/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 决定，模型与数据集必须一致，这里统一下发。 diff --git a/utils/__init__.py b/utils/__init__.py index 2fc0559..008b1c8 100644 --- a/utils/__init__.py +++ b/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", diff --git a/utils/device.py b/utils/device.py index fa3f59c..4fb3beb 100644 --- a/utils/device.py +++ b/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"