96 lines
4 KiB
Python
96 lines
4 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Benchmark CPU (float32) vs CPU (int8 quantized) vs VitisAI (int8 quantized static).
|
|
"""
|
|
|
|
import os
|
|
import time
|
|
import numpy as np
|
|
from pathlib import Path
|
|
|
|
FLOAT32_MODEL = Path.home() / ".cache/breadsearch/models/model.onnx"
|
|
QUANT_MODEL = Path.home() / ".cache/breadsearch/models/model_quantized.onnx"
|
|
STATIC_MODEL = Path.home() / ".cache/breadsearch/models/model_quantized_static.onnx"
|
|
CACHE_DIR = Path.home() / ".cache/breadsearch/npu/nomic-quantized-static"
|
|
VAIP_CONFIG = Path.home() / ".config/breadsearch/vaip_config.json"
|
|
RYZEN_AI_LIB = Path.home() / ".local/share/ryzen-ai-1.7.1/lib"
|
|
|
|
os.environ["RYZEN_AI_INSTALLATION_PATH"] = str(RYZEN_AI_LIB)
|
|
os.environ["LD_LIBRARY_PATH"] = str(RYZEN_AI_LIB) + ":" + os.environ.get("LD_LIBRARY_PATH", "")
|
|
|
|
import onnxruntime as ort
|
|
|
|
WARMUP = 2
|
|
RUNS = 10
|
|
SEQ_FLOAT = 512 # dynamic model accepts any seq len
|
|
SEQ_STATIC = 512 # static model locked to this
|
|
|
|
def make_input(seq_len: int):
|
|
ids = np.ones((1, seq_len), dtype=np.int64)
|
|
mask = np.ones((1, seq_len), dtype=np.int64)
|
|
types = np.zeros((1, seq_len), dtype=np.int64)
|
|
return {"input_ids": ids, "token_type_ids": types, "attention_mask": mask}
|
|
|
|
def time_session(sess: ort.InferenceSession, feed: dict, n: int) -> list[float]:
|
|
times = []
|
|
for _ in range(n):
|
|
t0 = time.perf_counter()
|
|
sess.run(None, feed)
|
|
times.append(time.perf_counter() - t0)
|
|
return times
|
|
|
|
def stats(times):
|
|
arr = np.array(times)
|
|
return arr.mean(), arr.min(), arr.max()
|
|
|
|
print("=" * 60)
|
|
print("BENCHMARK: nomic-embed-text-v1.5 embedding speed")
|
|
print("=" * 60)
|
|
|
|
# ── 1. float32 CPU ────────────────────────────────────────────
|
|
print("\n[1] float32 model — CPU EP")
|
|
sess = ort.InferenceSession(str(FLOAT32_MODEL), providers=["CPUExecutionProvider"])
|
|
feed = make_input(SEQ_FLOAT)
|
|
for _ in range(WARMUP): sess.run(None, feed)
|
|
times = time_session(sess, feed, RUNS)
|
|
mean, lo, hi = stats(times)
|
|
print(f" seq={SEQ_FLOAT} mean={mean*1000:.0f}ms min={lo*1000:.0f}ms max={hi*1000:.0f}ms ({RUNS} runs)")
|
|
|
|
# ── 2. int8 quantized CPU (dynamic) ───────────────────────────
|
|
print("\n[2] int8 quantized model (dynamic shapes) — CPU EP")
|
|
sess = ort.InferenceSession(str(QUANT_MODEL), providers=["CPUExecutionProvider"])
|
|
feed = make_input(SEQ_FLOAT)
|
|
for _ in range(WARMUP): sess.run(None, feed)
|
|
times = time_session(sess, feed, RUNS)
|
|
mean2, lo2, hi2 = stats(times)
|
|
print(f" seq={SEQ_FLOAT} mean={mean2*1000:.0f}ms min={lo2*1000:.0f}ms max={hi2*1000:.0f}ms ({RUNS} runs)")
|
|
print(f" Speedup vs float32: {mean/mean2:.2f}x")
|
|
|
|
# ── 3. int8 quantized static — VitisAI EP ─────────────────────
|
|
print("\n[3] int8 quantized model (static shapes) — VitisAI EP (NPU+CPU)")
|
|
providers = [
|
|
("VitisAIExecutionProvider", {
|
|
"config_file": str(VAIP_CONFIG),
|
|
"cacheDir": str(CACHE_DIR),
|
|
"cacheKey": "nomic-quantized-static",
|
|
}),
|
|
"CPUExecutionProvider",
|
|
]
|
|
print(" Loading session (should be fast — already compiled)...")
|
|
t_load = time.perf_counter()
|
|
sess_npu = ort.InferenceSession(str(STATIC_MODEL), providers=providers)
|
|
print(f" Load time: {time.perf_counter()-t_load:.1f}s")
|
|
feed_static = make_input(SEQ_STATIC)
|
|
for _ in range(WARMUP): sess_npu.run(None, feed_static)
|
|
times_npu = time_session(sess_npu, feed_static, RUNS)
|
|
mean3, lo3, hi3 = stats(times_npu)
|
|
print(f" seq={SEQ_STATIC} mean={mean3*1000:.0f}ms min={lo3*1000:.0f}ms max={hi3*1000:.0f}ms ({RUNS} runs)")
|
|
print(f" Speedup vs float32: {mean/mean3:.2f}x")
|
|
print(f" Speedup vs int8 CPU: {mean2/mean3:.2f}x")
|
|
|
|
print("\n" + "=" * 60)
|
|
print("SUMMARY")
|
|
print(f" float32 CPU : {mean*1000:.0f}ms/inference")
|
|
print(f" int8 CPU : {mean2*1000:.0f}ms/inference ({mean/mean2:.2f}x)")
|
|
print(f" int8 VitisAI : {mean3*1000:.0f}ms/inference ({mean/mean3:.2f}x)")
|
|
print("=" * 60)
|