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
+3
-3
@@ -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,
|
||||
|
||||
Reference in New Issue
Block a user