#!/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 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="默认有 GPU 用 cuda，否则 cpu") 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 == "cuda"): 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 is_cuda = device.type == "cuda" for _ in range(args.warmup): run_once(model, images, args) if is_cuda: torch.cuda.synchronize(device) torch.cuda.reset_peak_memory_stats(device) start = time.perf_counter() for _ in range(args.iters): run_once(model, images, args) if is_cuda: torch.cuda.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 is_cuda: result.update( { "peak_allocated_mb": torch.cuda.max_memory_allocated(device) / 1024**2, "peak_reserved_mb": torch.cuda.max_memory_reserved(device) / 1024**2, "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" 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) 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() device = torch.device(args.device or ("cuda" if torch.cuda.is_available() else "cpu")) if device.type != "cuda": logger.warning("当前在 %s 上运行：没有显存数据，只能看耗时与 CPU 内存。", device) 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 | device %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())