File size: 3,733 Bytes
fd0b01f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 | # cpu latency benchmark for the exported onnx. RUN LOCALLY on the target
# laptop. gpu timings are not informative because deployment is cpu-only via
# the rust ort runtime. mirrors privacy-filter/scripts/05_benchmark.py.
import argparse
import json
from pathlib import Path
from transformers import AutoTokenizer
from cleanup.eval.latency import benchmark_latency, benchmark_realistic
def _resolve_model_path(run_dir: Path, which: str) -> Path:
if which == "int8":
int8_path = run_dir / "onnx" / "int8" / "model.onnx"
if not int8_path.exists():
raise FileNotFoundError(f"no int8 onnx at {int8_path}; run scripts/04_export.py")
return int8_path
fp32_path = run_dir / "onnx" / "model.onnx"
if not fp32_path.exists():
raise FileNotFoundError(f"no fp32 onnx at {fp32_path}; run scripts/04_export.py")
return fp32_path
def _load_realistic_texts(data_dir: Path, n: int) -> list[str]:
test_path = Path(data_dir) / "test.json"
if not test_path.exists():
return []
rows = json.loads(test_path.read_text())
if n < len(rows):
rows = rows[:n]
return [r["raw"] for r in rows]
def main() -> None:
parser = argparse.ArgumentParser()
parser.add_argument("--runs-dir", default="runs")
parser.add_argument("--run-id", required=True)
parser.add_argument("--data-dir", default="data/pairs")
parser.add_argument("--model", choices=["fp32", "int8"], default="fp32")
parser.add_argument("--threads", type=int, default=4)
parser.add_argument("--warmup", type=int, default=50)
parser.add_argument("--measure", type=int, default=500)
parser.add_argument("--realistic-samples", type=int, default=500)
args = parser.parse_args()
run_dir = Path(args.runs_dir) / args.run_id
model_path = _resolve_model_path(run_dir, args.model)
print(f"[bench] {args.model} model: {model_path}")
print(f"[bench] file size {model_path.stat().st_size / 1e6:.1f} MB")
tokenizer = AutoTokenizer.from_pretrained(run_dir / "model", use_fast=True)
print("[bench] fixed length sweep")
sweep = benchmark_latency(
onnx_path=model_path,
tokenizer=tokenizer,
warmup=args.warmup,
measure=args.measure,
intra_op_threads=args.threads,
)
realistic = None
texts = _load_realistic_texts(Path(args.data_dir), args.realistic_samples)
if texts:
print(f"[bench] realistic mix on {len(texts)} real test rows")
realistic = benchmark_realistic(
onnx_path=model_path,
tokenizer=tokenizer,
texts=texts,
intra_op_threads=args.threads,
)
out_path = run_dir / "latency_benchmark.json"
out_path.write_text(json.dumps(
{
"model": args.model,
"model_path": str(model_path),
"model_size_bytes": model_path.stat().st_size,
"intra_op_threads": args.threads,
"results_by_length": sweep,
"realistic_mix": realistic,
},
indent=2,
))
print(f"[bench] wrote {out_path}")
print()
print("length | p50 ms | p95 ms | p99 ms | mean ms")
for length, stats in sweep.items():
print(
f"{length:>6s} | {stats['p50_ms']:>6.2f} | {stats['p95_ms']:>6.2f} | "
f"{stats['p99_ms']:>6.2f} | {stats['mean_ms']:>7.2f}"
)
if realistic:
print()
print(
f"realistic mix ({realistic['samples']} rows, "
f"p50 length {realistic['token_length_p50']}): "
f"p50={realistic['p50_ms']:.2f}ms p95={realistic['p95_ms']:.2f}ms "
f"p99={realistic['p99_ms']:.2f}ms"
)
if __name__ == "__main__":
main()
|