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