Shaun ZhouandClaude Opus 5
141cb907f2
移除"训练必须有 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>
40 lines
1.1 KiB
Python
40 lines
1.1 KiB
Python
"""按配置构建 Dataset / DataLoader。"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from collections.abc import Mapping
|
|
from typing import Any
|
|
|
|
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:
|
|
return DATASETS.build(dict(cfg), **kwargs)
|
|
|
|
|
|
def build_dataloader(
|
|
dataset: Dataset,
|
|
batch_size: int = 8,
|
|
shuffle: bool = False,
|
|
num_workers: int = 4,
|
|
pin_memory: bool | None = None,
|
|
drop_last: bool = False,
|
|
) -> DataLoader:
|
|
"""`pin_memory` 缺省时只在 CUDA 上开启（CPU/MPS 上无收益且会告警）。"""
|
|
if pin_memory is None:
|
|
pin_memory = pin_memory_supported(resolve_device("auto"))
|
|
return DataLoader(
|
|
dataset,
|
|
batch_size=batch_size,
|
|
shuffle=shuffle,
|
|
num_workers=num_workers,
|
|
pin_memory=pin_memory,
|
|
drop_last=drop_last,
|
|
collate_fn=collate_fn,
|
|
persistent_workers=num_workers > 0,
|
|
)
|