Initial release
This commit is contained in:
civ
@@ -0,0 +1,269 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Validate and benchmark VAE ONNX model with TensorRT Execution Provider.
|
||||
|
||||
Compares CUDA EP (baseline) against TensorRT EP for latency and numerical
|
||||
accuracy. Caches TRT engines for subsequent runs.
|
||||
|
||||
Usage:
|
||||
python test_trt_vae.py --onnx models/onnx/vae_decoder.onnx
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import os
|
||||
import sys
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
|
||||
|
||||
def check_providers():
|
||||
"""Check which ONNX Runtime execution providers are available."""
|
||||
try:
|
||||
import onnxruntime as ort
|
||||
except ImportError:
|
||||
print("ERROR: onnxruntime is not installed.")
|
||||
print("Install with: pip install onnxruntime-gpu")
|
||||
sys.exit(1)
|
||||
|
||||
available = ort.get_available_providers()
|
||||
print(f"onnxruntime version: {ort.__version__}")
|
||||
print(f"Available providers: {available}")
|
||||
|
||||
has_cuda = "CUDAExecutionProvider" in available
|
||||
has_trt = "TensorrtExecutionProvider" in available
|
||||
|
||||
if not has_cuda:
|
||||
print("\n[WARN] CUDAExecutionProvider is NOT available.")
|
||||
print(" You likely have `onnxruntime` (CPU-only) instead of `onnxruntime-gpu`.")
|
||||
print(" Install with: pip install onnxruntime-gpu")
|
||||
print(" (You may need to uninstall onnxruntime first)")
|
||||
|
||||
# Check which package is installed
|
||||
try:
|
||||
import importlib.metadata
|
||||
try:
|
||||
ver = importlib.metadata.version("onnxruntime-gpu")
|
||||
print(f" onnxruntime-gpu version: {ver}")
|
||||
except importlib.metadata.PackageNotFoundError:
|
||||
print(" onnxruntime-gpu: NOT installed")
|
||||
try:
|
||||
ver = importlib.metadata.version("onnxruntime")
|
||||
print(f" onnxruntime (CPU): {ver}")
|
||||
except importlib.metadata.PackageNotFoundError:
|
||||
pass
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
return has_cuda, has_trt
|
||||
|
||||
|
||||
def benchmark_session(sess, input_data, warmup=3, iterations=20):
|
||||
"""Benchmark an ONNX Runtime session.
|
||||
|
||||
Args:
|
||||
sess: ONNX Runtime InferenceSession.
|
||||
input_data: Dict of input name → numpy array.
|
||||
warmup: Number of warmup iterations.
|
||||
iterations: Number of timed iterations.
|
||||
|
||||
Returns:
|
||||
Tuple of (output_array, mean_latency_ms, std_latency_ms).
|
||||
"""
|
||||
# Warmup
|
||||
for _ in range(warmup):
|
||||
output = sess.run(None, input_data)
|
||||
|
||||
# Timed runs
|
||||
latencies = []
|
||||
for _ in range(iterations):
|
||||
t0 = time.perf_counter()
|
||||
output = sess.run(None, input_data)
|
||||
latencies.append((time.perf_counter() - t0) * 1000)
|
||||
|
||||
latencies = np.array(latencies)
|
||||
return output[0], latencies.mean(), latencies.std()
|
||||
|
||||
|
||||
def run_benchmark(onnx_path: str, trt_cache_dir: str):
|
||||
"""Run the full benchmark comparing CUDA EP vs TRT EP.
|
||||
|
||||
Args:
|
||||
onnx_path: Path to the ONNX model file.
|
||||
trt_cache_dir: Directory to cache TRT engines.
|
||||
"""
|
||||
import onnxruntime as ort
|
||||
|
||||
has_cuda, has_trt = check_providers()
|
||||
|
||||
if not has_cuda:
|
||||
print("\nCannot run GPU benchmarks without CUDAExecutionProvider.")
|
||||
print("Falling back to CPU-only test...")
|
||||
run_cpu_test(onnx_path)
|
||||
return
|
||||
|
||||
# Test input: 10 seconds of audio (250 latent frames)
|
||||
test_latents = np.random.randn(1, 64, 250).astype(np.float32)
|
||||
input_data = {"latents": test_latents}
|
||||
print(f"\nTest input shape: {test_latents.shape}")
|
||||
print(f"Expected output: [1, 2, {250 * 1920}] samples")
|
||||
|
||||
# ── CUDA EP benchmark ──
|
||||
print("\n" + "=" * 60)
|
||||
print("CUDA EP Benchmark")
|
||||
print("=" * 60)
|
||||
|
||||
cuda_opts = ort.SessionOptions()
|
||||
cuda_opts.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
|
||||
|
||||
try:
|
||||
cuda_sess = ort.InferenceSession(
|
||||
onnx_path,
|
||||
sess_options=cuda_opts,
|
||||
providers=["CUDAExecutionProvider", "CPUExecutionProvider"],
|
||||
)
|
||||
except Exception as e:
|
||||
print(f"Failed to create CUDA session: {e}")
|
||||
print("Falling back to CPU-only test...")
|
||||
run_cpu_test(onnx_path)
|
||||
return
|
||||
|
||||
cuda_output, cuda_mean, cuda_std = benchmark_session(cuda_sess, input_data)
|
||||
print(f"Output shape: {cuda_output.shape}")
|
||||
print(f"Latency: {cuda_mean:.2f} ± {cuda_std:.2f} ms")
|
||||
|
||||
# ── TensorRT EP benchmark ──
|
||||
if not has_trt:
|
||||
print("\n" + "=" * 60)
|
||||
print("TensorRT EP: NOT AVAILABLE")
|
||||
print("=" * 60)
|
||||
print("TensorrtExecutionProvider is not available in this onnxruntime build.")
|
||||
print("To enable TRT:")
|
||||
print(" 1. Install onnxruntime-gpu with TRT support")
|
||||
print(" 2. Ensure TensorRT libraries are on PATH")
|
||||
print("\nSkipping TRT benchmark. CUDA EP results above are the baseline.")
|
||||
return
|
||||
|
||||
print("\n" + "=" * 60)
|
||||
print("TensorRT EP Benchmark")
|
||||
print("=" * 60)
|
||||
|
||||
os.makedirs(trt_cache_dir, exist_ok=True)
|
||||
|
||||
trt_opts = ort.SessionOptions()
|
||||
trt_opts.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
|
||||
|
||||
trt_provider_options = {
|
||||
"trt_fp16_enable": True,
|
||||
"trt_engine_cache_enable": True,
|
||||
"trt_engine_cache_path": trt_cache_dir,
|
||||
"trt_max_workspace_size": str(8 * 1024 * 1024 * 1024), # 8GB
|
||||
}
|
||||
|
||||
print("Building TRT engine (first run may take minutes)...")
|
||||
t0 = time.time()
|
||||
try:
|
||||
trt_sess = ort.InferenceSession(
|
||||
onnx_path,
|
||||
sess_options=trt_opts,
|
||||
providers=[
|
||||
("TensorrtExecutionProvider", trt_provider_options),
|
||||
"CUDAExecutionProvider",
|
||||
"CPUExecutionProvider",
|
||||
],
|
||||
)
|
||||
except Exception as e:
|
||||
print(f"Failed to create TRT session: {e}")
|
||||
print("TRT EP may not be properly configured. Skipping TRT benchmark.")
|
||||
return
|
||||
|
||||
engine_time = time.time() - t0
|
||||
print(f"TRT engine ready in {engine_time:.1f}s")
|
||||
|
||||
trt_output, trt_mean, trt_std = benchmark_session(trt_sess, input_data)
|
||||
print(f"Output shape: {trt_output.shape}")
|
||||
print(f"Latency: {trt_mean:.2f} ± {trt_std:.2f} ms")
|
||||
|
||||
# ── Comparison ──
|
||||
print("\n" + "=" * 60)
|
||||
print("Comparison: CUDA EP vs TensorRT EP")
|
||||
print("=" * 60)
|
||||
|
||||
abs_diff = np.abs(cuda_output - trt_output)
|
||||
max_diff = abs_diff.max()
|
||||
mean_diff = abs_diff.mean()
|
||||
|
||||
print(f"CUDA EP latency: {cuda_mean:.2f} ± {cuda_std:.2f} ms")
|
||||
print(f"TRT EP latency: {trt_mean:.2f} ± {trt_std:.2f} ms")
|
||||
print(f"Speedup: {cuda_mean / trt_mean:.2f}x")
|
||||
print(f"Max abs diff: {max_diff:.6e}")
|
||||
print(f"Mean abs diff: {mean_diff:.6e}")
|
||||
|
||||
if max_diff < 0.05:
|
||||
print("\n[PASS] Numerical accuracy: GOOD (fp16 rounding is expected)")
|
||||
elif max_diff < 0.5:
|
||||
print("\n[WARN] Numerical accuracy: ACCEPTABLE (fp16 precision loss)")
|
||||
else:
|
||||
print(f"\n[FAIL] Numerical accuracy: POOR (max diff = {max_diff:.4f})")
|
||||
print(" This may indicate a TRT conversion issue.")
|
||||
|
||||
|
||||
def run_cpu_test(onnx_path: str):
|
||||
"""Fallback: run a basic CPU test to verify the ONNX model loads."""
|
||||
import onnxruntime as ort
|
||||
|
||||
print("\n" + "=" * 60)
|
||||
print("CPU-only Test (fallback)")
|
||||
print("=" * 60)
|
||||
|
||||
test_latents = np.random.randn(1, 64, 50).astype(np.float32)
|
||||
input_data = {"latents": test_latents}
|
||||
|
||||
sess = ort.InferenceSession(onnx_path, providers=["CPUExecutionProvider"])
|
||||
output = sess.run(None, input_data)
|
||||
|
||||
print(f"Input shape: {test_latents.shape}")
|
||||
print(f"Output shape: {output[0].shape}")
|
||||
print(f"Output range: [{output[0].min():.4f}, {output[0].max():.4f}]")
|
||||
print("[PASS] Model loads and runs on CPU successfully.")
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Validate and benchmark VAE ONNX model with TensorRT EP"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--onnx",
|
||||
type=str,
|
||||
required=True,
|
||||
help="Path to the exported ONNX file",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--trt-cache",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Directory for TRT engine cache (default: alongside ONNX file)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--iterations",
|
||||
type=int,
|
||||
default=20,
|
||||
help="Number of benchmark iterations (default: 20)",
|
||||
)
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
if not os.path.isfile(args.onnx):
|
||||
print(f"ERROR: ONNX file not found: {args.onnx}")
|
||||
sys.exit(1)
|
||||
|
||||
trt_cache = args.trt_cache
|
||||
if trt_cache is None:
|
||||
trt_cache = os.path.join(os.path.dirname(args.onnx), "trt_cache")
|
||||
|
||||
run_benchmark(args.onnx, trt_cache)
|
||||
print("\nDone!")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user