Shaun ZhouandClaude Opus 5
4b1147ed8f
- convert_annotations: labelme / VOC -> YOLO txt，顺带生成 data.yaml
- split_dataset: 切分成 images/{train,val} 与 labels/{train,val}，默认建软链不复制数据
- export_onnx: 导出部署用 ONNX，CPU 可跑，缺依赖时提示 `uv sync --group export`
- profile_inference: torch.profiler 测单帧延迟/FPS、CUDA 显存峰值与算子级明细
- check_env: Python / PyTorch / Nvidia GPU 自检
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
205 lines
7.9 KiB
Python
Executable File
205 lines
7.9 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
|
||
|
||
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())
|