2026-03-18 17:23:23 +08:00

222 lines
7.2 KiB
Python

#!/usr/bin/env python3
"""
BPU Model Performance Benchmark Tool
Runs hrt_model_exec perf for each model at thread counts 1,2,3,4
"""
import argparse
import json
import re
import subprocess
import sys
from pathlib import Path
HRT_MODEL_EXEC = "hrt_model_exec"
THREAD_COUNTS = [1, 2, 3, 4]
def run_perf(model_path: str, thread_num: int, frame_count: int = 200) -> dict:
"""Run hrt_model_exec perf and return parsed results."""
cmd = [
HRT_MODEL_EXEC, "perf",
"--model_file", model_path,
"--thread_num", str(thread_num),
"--frame_count", str(frame_count),
]
result = subprocess.run(cmd, capture_output=True, text=True)
output = result.stdout + result.stderr
perf = {
"thread_num": thread_num,
"frame_count": frame_count,
"run_time_ms": None,
"total_latency_ms": None,
"avg_latency_ms": None,
"fps": None,
"raw_output": output,
"returncode": result.returncode,
}
m = re.search(r"Program run time:\s*([\d.]+)\s*ms", output)
if m:
perf["run_time_ms"] = float(m.group(1))
m = re.search(r"Frame totally latency is:\s*([\d.]+)\s*ms", output)
if m:
perf["total_latency_ms"] = float(m.group(1))
m = re.search(r"Average\s+latency\s+is:\s*([\d.]+)\s*ms", output)
if m:
perf["avg_latency_ms"] = float(m.group(1))
m = re.search(r"Frame\s+rate\s+is:\s*([\d.]+)\s*FPS", output)
if m:
perf["fps"] = float(m.group(1))
return perf
def print_table(results: list):
"""Print results as a human-readable table."""
# 动态计算 model 列宽
max_name_len = max(len(e["model_name"]) for e in results)
col_widths = [max(max_name_len, 10), 9, 16, 10]
headers = ["Model", "Threads", "Avg Latency(ms)", "FPS"]
sep = "+" + "+".join("-" * (w + 2) for w in col_widths) + "+"
header_row = "|" + "|".join(
f" {h:<{w}} " for h, w in zip(headers, col_widths)
) + "|"
for entry in results:
name = entry["model_name"]
# 每个模型单独打印表头
print(sep)
print(header_row)
print(sep)
for p in entry["perf_results"]:
avg = f"{p['avg_latency_ms']:.3f}" if p["avg_latency_ms"] is not None else "N/A"
fps = f"{p['fps']:.2f}" if p["fps"] is not None else "N/A"
row = [name, str(p["thread_num"]), avg, fps]
print("|" + "|".join(
f" {v:<{w}} " for v, w in zip(row, col_widths)
) + "|")
name = ""
print(sep)
def validate_config(config: dict, input_dir: Path) -> list:
"""Print and validate task.json, return normalized model list. Exit on error."""
print("=" * 60)
print("task.json content:")
print(json.dumps(config, indent=2, ensure_ascii=False))
print("=" * 60)
errors = []
# --- frame_count ---
if "frame_count" in config:
if not isinstance(config["frame_count"], int) or config["frame_count"] <= 0:
errors.append(" [frame_count] must be a positive integer")
# --- 判断格式 ---
has_single = "model_relative_path" in config
has_multi = "models" in config
if not has_single and not has_multi:
errors.append(" missing required field: 'model_relative_path' or 'models'")
for e in errors:
print(f"[ERROR] {e}", file=sys.stderr)
sys.exit(1)
if has_single and has_multi:
errors.append(" ambiguous: both 'model_relative_path' and 'models' are present, use one")
models = []
if has_single:
rel = config["model_relative_path"]
if not isinstance(rel, str) or not rel.strip():
errors.append(" [model_relative_path] must be a non-empty string")
elif not rel.endswith((".hbm", ".bin")):
errors.append(f" [model_relative_path] unsupported extension: '{rel}' (expected .hbm or .bin)")
else:
models = [{"name": Path(rel).name, "path": rel}]
if has_multi:
if not isinstance(config["models"], list) or len(config["models"]) == 0:
errors.append(" [models] must be a non-empty list")
else:
for i, m in enumerate(config["models"]):
prefix = f" [models[{i}]]"
if not isinstance(m, dict):
errors.append(f"{prefix} each entry must be an object")
continue
if "path" not in m:
errors.append(f"{prefix} missing required field 'path'")
elif not isinstance(m["path"], str) or not m["path"].strip():
errors.append(f"{prefix} 'path' must be a non-empty string")
elif not m["path"].endswith((".hbm", ".bin")):
errors.append(f"{prefix} unsupported extension: '{m['path']}' (expected .hbm or .bin)")
else:
models.append({"name": m.get("name", Path(m["path"]).name), "path": m["path"]})
if errors:
for e in errors:
print(f"[ERROR] {e}", file=sys.stderr)
sys.exit(1)
# --- 文件存在性检查 ---
missing = []
for m in models:
full = input_dir / m["path"]
if not full.exists():
missing.append(f" model file not found: {full}")
if missing:
for e in missing:
print(f"[ERROR] {e}", file=sys.stderr)
sys.exit(1)
print(f"[OK] {len(models)} model(s) validated, frame_count={config.get('frame_count', 200)}\n")
return models
def main():
parser = argparse.ArgumentParser(description="BPU model perf benchmark")
parser.add_argument("--input", required=True, help="Input JSON file")
parser.add_argument("--output", required=True, help="Output JSON file")
args = parser.parse_args()
if not Path(args.input).exists():
print(f"[ERROR] input file not found: {args.input}", file=sys.stderr)
sys.exit(1)
with open(args.input) as f:
config = json.load(f)
input_dir = Path(args.input).parent
frame_count = config.get("frame_count", 200)
models = validate_config(config, input_dir)
if not models:
print("[ERROR] no models specified in input JSON", file=sys.stderr)
sys.exit(1)
output_results = []
for model in models:
rel_path = model["path"]
model_path = str(input_dir / rel_path)
model_name = model.get("name", Path(rel_path).name)
print(f"\n[Benchmarking] {model_name} ({model_path})")
perf_results = []
for t in THREAD_COUNTS:
print(f" thread_num={t} ...", end=" ", flush=True)
p = run_perf(model_path, t, frame_count)
perf_results.append(p)
if p["fps"] is not None:
print(f"FPS={p['fps']:.2f} avg_latency={p['avg_latency_ms']:.3f}ms")
else:
print("FAILED (check raw_output in result JSON)")
output_results.append({
"model_name": model_name,
"model_path": model_path,
"perf_results": perf_results,
})
print("\n" + "=" * 80)
print_table(output_results)
output_path = Path(args.output)
output_path.parent.mkdir(parents=True, exist_ok=True)
with open(output_path, "w") as f:
json.dump(output_results, f, indent=2)
print(f"\nResults saved to: {args.output}")
if __name__ == "__main__":
main()