mirror of
https://github.com/firestar5683/StarPilot.git
synced 2026-07-22 17:52:07 +08:00
firestar5683
69 lines
2.0 KiB
Python
69 lines
2.0 KiB
Python
#!/usr/bin/env python3
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
import statistics
|
|
import time
|
|
|
|
from pathlib import Path
|
|
|
|
import cv2
|
|
import numpy as np
|
|
|
|
|
|
def parse_args() -> argparse.Namespace:
|
|
parser = argparse.ArgumentParser(description="Benchmark fixed-shape ONNX forwards through OpenCV DNN.")
|
|
parser.add_argument("models", nargs="+", type=Path, help="ONNX model paths to benchmark.")
|
|
parser.add_argument("--input-size", type=int, required=True, help="Square model input size.")
|
|
parser.add_argument("--iterations", type=int, default=50)
|
|
parser.add_argument("--warmup", type=int, default=5)
|
|
return parser.parse_args()
|
|
|
|
|
|
def percentile(values: list[float], quantile: float) -> float:
|
|
ordered = sorted(values)
|
|
index = min(round((len(ordered) - 1) * quantile), len(ordered) - 1)
|
|
return ordered[index]
|
|
|
|
|
|
def benchmark(path: Path, input_size: int, iterations: int, warmup: int) -> dict[str, object]:
|
|
net = cv2.dnn.readNetFromONNX(str(path))
|
|
blob = np.zeros((1, 3, input_size, input_size), dtype=np.float32)
|
|
net.setInput(blob)
|
|
output = net.forward()
|
|
for _ in range(max(warmup - 1, 0)):
|
|
net.setInput(blob)
|
|
output = net.forward()
|
|
|
|
durations_ms = []
|
|
for _ in range(max(iterations, 1)):
|
|
net.setInput(blob)
|
|
started_at = time.perf_counter()
|
|
output = net.forward()
|
|
durations_ms.append((time.perf_counter() - started_at) * 1000.0)
|
|
|
|
return {
|
|
"model": str(path.resolve()),
|
|
"input_size": input_size,
|
|
"output_shape": list(output.shape),
|
|
"iterations": len(durations_ms),
|
|
"median_ms": round(statistics.median(durations_ms), 3),
|
|
"p95_ms": round(percentile(durations_ms, 0.95), 3),
|
|
"mean_ms": round(statistics.mean(durations_ms), 3),
|
|
}
|
|
|
|
|
|
def main() -> int:
|
|
args = parse_args()
|
|
results = [
|
|
benchmark(path.expanduser().resolve(), args.input_size, args.iterations, args.warmup)
|
|
for path in args.models
|
|
]
|
|
print(json.dumps(results, indent=2))
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|