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:
co-authored by
Claude Opus 5
parent
4fc0fa5353
commit
141cb907f2
@@ -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 决定,模型与数据集必须一致,这里统一下发。
|
||||
|
||||
Reference in New Issue
Block a user