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
co-authored by
Claude Opus 5
Claude Opus 5
parent
4fc0fa5353
commit
141cb907f2
@@ -31,6 +31,12 @@ from torch.profiler import ProfilerActivity, profile, record_function # noqa: E
|
||||
from datasets import load_classes # noqa: E402
|
||||
from models import build_model, decode_heatmap # noqa: E402
|
||||
from utils import load_config, setup_logging # noqa: E402
|
||||
from utils.device import ( # noqa: E402
|
||||
DeviceNotAvailableError,
|
||||
PeakMemory,
|
||||
resolve_device,
|
||||
synchronize,
|
||||
)
|
||||
|
||||
logger = setup_logging()
|
||||
|
||||
@@ -45,7 +51,7 @@ def parse_args() -> argparse.Namespace:
|
||||
parser.add_argument("--image-size", type=int, nargs=2, default=None, help="默认取配置的 image_size")
|
||||
parser.add_argument("--iters", type=int, default=50, help="计时轮数")
|
||||
parser.add_argument("--warmup", type=int, default=10, help="预热轮数（不计入统计）")
|
||||
parser.add_argument("--device", default=None, help="默认有 GPU 用 cuda，否则 cpu")
|
||||
parser.add_argument("--device", default=None, help="auto / cuda / mps / cpu，默认 auto")
|
||||
parser.add_argument("--amp", action="store_true", help="用 float16 autocast 推理")
|
||||
parser.add_argument("--postprocess", action="store_true", help="把 heatmap 解码也计入")
|
||||
parser.add_argument("--rows", type=int, default=15, help="算子排行显示条数")
|
||||
@@ -73,7 +79,7 @@ def build(args: argparse.Namespace, device: torch.device) -> torch.nn.Module:
|
||||
@torch.no_grad()
|
||||
def run_once(model: torch.nn.Module, images: torch.Tensor, args: argparse.Namespace) -> None:
|
||||
device_type = images.device.type
|
||||
with torch.autocast(device_type, dtype=torch.float16, enabled=args.amp and device_type == "cuda"):
|
||||
with torch.autocast(device_type, dtype=torch.float16, enabled=args.amp and device_type != "cpu"):
|
||||
with record_function("forward"):
|
||||
outputs = model(images)
|
||||
if args.postprocess:
|
||||
@@ -89,19 +95,17 @@ def run_once(model: torch.nn.Module, images: torch.Tensor, args: argparse.Namesp
|
||||
def measure(model: torch.nn.Module, images: torch.Tensor, args: argparse.Namespace) -> dict:
|
||||
"""预热后计时，并统计 CUDA 显存峰值。"""
|
||||
device = images.device
|
||||
is_cuda = device.type == "cuda"
|
||||
peak = PeakMemory(device)
|
||||
|
||||
for _ in range(args.warmup):
|
||||
run_once(model, images, args)
|
||||
if is_cuda:
|
||||
torch.cuda.synchronize(device)
|
||||
torch.cuda.reset_peak_memory_stats(device)
|
||||
synchronize(device)
|
||||
peak.reset()
|
||||
|
||||
start = time.perf_counter()
|
||||
for _ in range(args.iters):
|
||||
run_once(model, images, args)
|
||||
if is_cuda:
|
||||
torch.cuda.synchronize(device)
|
||||
synchronize(device)
|
||||
elapsed = time.perf_counter() - start
|
||||
|
||||
batch = images.shape[0]
|
||||
@@ -110,11 +114,12 @@ def measure(model: torch.nn.Module, images: torch.Tensor, args: argparse.Namespa
|
||||
"latency_ms": elapsed / args.iters * 1000,
|
||||
"throughput_fps": args.iters * batch / elapsed,
|
||||
}
|
||||
if is_cuda:
|
||||
if device.type != "cpu":
|
||||
result.update(
|
||||
{
|
||||
"peak_allocated_mb": torch.cuda.max_memory_allocated(device) / 1024**2,
|
||||
"peak_reserved_mb": torch.cuda.max_memory_reserved(device) / 1024**2,
|
||||
"peak_allocated_mb": peak.peak_mb,
|
||||
"peak_reserved_mb": peak.reserved_mb,
|
||||
"peak_is_exact": peak.exact,
|
||||
"weights_mb": sum(p.numel() * p.element_size() for p in model.parameters()) / 1024**2,
|
||||
}
|
||||
)
|
||||
@@ -124,13 +129,13 @@ def measure(model: torch.nn.Module, images: torch.Tensor, args: argparse.Namespa
|
||||
def profile_once(model: torch.nn.Module, images: torch.Tensor, args: argparse.Namespace) -> None:
|
||||
"""算子级 profiling：按显存占用与耗时各排一次序。"""
|
||||
is_cuda = images.device.type == "cuda"
|
||||
# torch.profiler 没有 MPS 的 activity，Metal 侧的 kernel 抓不到，只能看 CPU 侧下发耗时
|
||||
activities = [ProfilerActivity.CPU] + ([ProfilerActivity.CUDA] if is_cuda else [])
|
||||
|
||||
with profile(activities=activities, profile_memory=True, record_shapes=True) as prof:
|
||||
for _ in range(max(args.iters // 5, 1)):
|
||||
run_once(model, images, args)
|
||||
if is_cuda:
|
||||
torch.cuda.synchronize(images.device)
|
||||
synchronize(images.device)
|
||||
|
||||
memory_key = "self_cuda_memory_usage" if is_cuda else "self_cpu_memory_usage"
|
||||
time_key = "self_cuda_time_total" if is_cuda else "self_cpu_time_total"
|
||||
@@ -148,9 +153,18 @@ def profile_once(model: torch.nn.Module, images: torch.Tensor, args: argparse.Na
|
||||
|
||||
def main() -> int:
|
||||
args = parse_args()
|
||||
device = torch.device(args.device or ("cuda" if torch.cuda.is_available() else "cpu"))
|
||||
if device.type != "cuda":
|
||||
logger.warning("当前在 %s 上运行：没有显存数据，只能看耗时与 CPU 内存。", device)
|
||||
try:
|
||||
device = resolve_device(args.device or "auto")
|
||||
except DeviceNotAvailableError as exc:
|
||||
logger.error("%s", exc)
|
||||
return 1
|
||||
if device.type == "cpu":
|
||||
logger.warning("当前在 CPU 上运行：没有显存数据，只能看耗时与 CPU 内存。")
|
||||
elif device.type == "mps":
|
||||
logger.warning(
|
||||
"MPS 上 torch.profiler 抓不到 Metal kernel（没有对应的 ProfilerActivity），"
|
||||
"算子明细只反映 CPU 侧下发耗时；显存用 torch.mps 采样，峰值是下界。"
|
||||
)
|
||||
|
||||
model = build(args, device)
|
||||
cfg = model.cfg
|
||||
@@ -158,7 +172,7 @@ def main() -> int:
|
||||
channels = 1 if cfg.get("gray") else 3
|
||||
|
||||
logger.info(
|
||||
"模型 %s | 参数量 %.2fM | 输出 stride %d | 输入 %dx%dx%d | device %s | amp=%s",
|
||||
"模型 %s | 参数量 %.2fM | 输出 stride %d | 输入 %dx%dx%d | 设备 %s | amp=%s",
|
||||
cfg["model"]["name"],
|
||||
sum(p.numel() for p in model.parameters()) / 1e6,
|
||||
model.out_stride,
|
||||
|
||||
Reference in New Issue
Block a user