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>
219 lines
8.3 KiB
Python
Executable File
219 lines
8.3 KiB
Python
Executable File
#!/usr/bin/env python
|
||
"""推理性能与显存分析（torch.profiler）。
|
||
|
||
部署在实时设备上，显存占用和单帧耗时是硬指标，因此这个脚本在推理路径上跑
|
||
`torch.profiler`，给出算子级的显存与耗时排行，外加 CUDA 显存峰值统计。
|
||
|
||
# 用配置里的输入尺寸，batch=1 跑 50 次
|
||
python scripts/profile_inference.py --config configs/train/baseline.yaml \
|
||
--checkpoint runs/baseline/best.pth
|
||
|
||
# 扫不同 batch，并导出 chrome trace（chrome://tracing 或 https://ui.perfetto.dev 打开）
|
||
python scripts/profile_inference.py --config configs/train/baseline.yaml \
|
||
--batch-size 1 4 8 --trace runs/baseline/trace.json
|
||
|
||
无 GPU 时自动退回 CPU（只有耗时与 CPU 内存，没有显存数据）。
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import argparse
|
||
import json
|
||
import sys
|
||
import time
|
||
from pathlib import Path
|
||
|
||
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
|
||
|
||
import torch # noqa: E402
|
||
from torch.profiler import ProfilerActivity, profile, record_function # noqa: E402
|
||
|
||
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()
|
||
|
||
|
||
def parse_args() -> argparse.Namespace:
|
||
parser = argparse.ArgumentParser(
|
||
description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter
|
||
)
|
||
parser.add_argument("--config", required=True)
|
||
parser.add_argument("--checkpoint", default=None, help="不给则用随机初始化的权重，只看结构开销")
|
||
parser.add_argument("--batch-size", type=int, nargs="+", default=[1])
|
||
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="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="算子排行显示条数")
|
||
parser.add_argument("--trace", default=None, help="导出 chrome trace 的路径")
|
||
parser.add_argument("--out", default=None, help="把汇总结果写入 json")
|
||
return parser.parse_args()
|
||
|
||
|
||
def build(args: argparse.Namespace, device: torch.device) -> torch.nn.Module:
|
||
cfg = load_config(args.config)
|
||
model_cfg = dict(cfg["model"])
|
||
model_cfg.setdefault("num_classes", len(load_classes(cfg["classes"])))
|
||
model_cfg.setdefault("label_mode", str(cfg.get("label_mode", "multilabel")))
|
||
model = build_model(model_cfg)
|
||
|
||
if args.checkpoint:
|
||
from engine import load_checkpoint
|
||
|
||
load_checkpoint(args.checkpoint, model, map_location="cpu")
|
||
model.to(device).eval()
|
||
model.cfg = cfg
|
||
return model
|
||
|
||
|
||
@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 != "cpu"):
|
||
with record_function("forward"):
|
||
outputs = model(images)
|
||
if args.postprocess:
|
||
with record_function("postprocess"):
|
||
decode_heatmap(
|
||
outputs["heatmap"].float(),
|
||
outputs["wh"].float(),
|
||
outputs["offset"].float(),
|
||
stride=model.out_stride,
|
||
)
|
||
|
||
|
||
def measure(model: torch.nn.Module, images: torch.Tensor, args: argparse.Namespace) -> dict:
|
||
"""预热后计时，并统计 CUDA 显存峰值。"""
|
||
device = images.device
|
||
peak = PeakMemory(device)
|
||
|
||
for _ in range(args.warmup):
|
||
run_once(model, images, args)
|
||
synchronize(device)
|
||
peak.reset()
|
||
|
||
start = time.perf_counter()
|
||
for _ in range(args.iters):
|
||
run_once(model, images, args)
|
||
synchronize(device)
|
||
elapsed = time.perf_counter() - start
|
||
|
||
batch = images.shape[0]
|
||
result = {
|
||
"batch_size": batch,
|
||
"latency_ms": elapsed / args.iters * 1000,
|
||
"throughput_fps": args.iters * batch / elapsed,
|
||
}
|
||
if device.type != "cpu":
|
||
result.update(
|
||
{
|
||
"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,
|
||
}
|
||
)
|
||
return result
|
||
|
||
|
||
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)
|
||
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"
|
||
print(f"\n===== 显存占用 Top {args.rows}（按 {memory_key}）=====")
|
||
print(prof.key_averages().table(sort_by=memory_key, row_limit=args.rows))
|
||
print(f"\n===== 耗时 Top {args.rows}（按 {time_key}）=====")
|
||
print(prof.key_averages().table(sort_by=time_key, row_limit=args.rows))
|
||
|
||
if args.trace:
|
||
trace_path = Path(args.trace)
|
||
trace_path.parent.mkdir(parents=True, exist_ok=True)
|
||
prof.export_chrome_trace(str(trace_path))
|
||
print(f"\nchrome trace 已导出: {trace_path}（用 chrome://tracing 或 ui.perfetto.dev 打开）")
|
||
|
||
|
||
def main() -> int:
|
||
args = parse_args()
|
||
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
|
||
height, width = args.image_size or cfg.get("image_size", [512, 512])
|
||
channels = 1 if cfg.get("gray") else 3
|
||
|
||
logger.info(
|
||
"模型 %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,
|
||
channels,
|
||
height,
|
||
width,
|
||
device,
|
||
args.amp,
|
||
)
|
||
|
||
results = []
|
||
for batch_size in args.batch_size:
|
||
images = torch.randn(batch_size, channels, int(height), int(width), device=device)
|
||
result = measure(model, images, args)
|
||
results.append(result)
|
||
memory = (
|
||
f" | 显存峰值 allocated {result['peak_allocated_mb']:.1f} MB"
|
||
f" / reserved {result['peak_reserved_mb']:.1f} MB"
|
||
if "peak_allocated_mb" in result
|
||
else ""
|
||
)
|
||
logger.info(
|
||
"batch=%-3d 单次 %.2f ms | %.1f FPS%s",
|
||
result["batch_size"],
|
||
result["latency_ms"],
|
||
result["throughput_fps"],
|
||
memory,
|
||
)
|
||
|
||
# 算子级明细只对最大的 batch 跑一次，避免输出过长。
|
||
images = torch.randn(max(args.batch_size), channels, int(height), int(width), device=device)
|
||
profile_once(model, images, args)
|
||
|
||
if args.out:
|
||
out_path = Path(args.out)
|
||
out_path.parent.mkdir(parents=True, exist_ok=True)
|
||
out_path.write_text(json.dumps(results, indent=2, ensure_ascii=False), encoding="utf-8")
|
||
logger.info("汇总已写入 %s", out_path)
|
||
return 0
|
||
|
||
|
||
if __name__ == "__main__":
|
||
raise SystemExit(main())
|