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>
This commit is contained in:
Shaun Zhou
2026-09-02 21:43:35 +08:00
co-authored by Claude Opus 5
parent 4fc0fa5353
commit 141cb907f2
14 changed files with 475 additions and 148 deletions
+16 -12
View File
@@ -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: