The Performance Battle comparison in TensorRT vs ONNX: The Truth

import tensorrt as trt
import pycuda.driver as cuda
import pycuda.autoinit
import numpy as np
import os
import time
from typing import List, Optional

class TensorRTInference:
    """Production-grade TensorRT inference wrapper with error handling and benchmarking."""

    def __init__(self, onnx_path: str, engine_path: str, precision: str = "fp16"):
        self.onnx_path = onnx_path
        self.engine_path = engine_path
        self.precision = precision.lower()
        self.logger = trt.Logger(trt.Logger.ERROR)
        self.engine: Optional[trt.ICudaEngine] = None
        self.context: Optional[trt.IExecutionContext] = None
        self.inputs: List[np.ndarray] = []
        self.outputs: List[np.ndarray] = []
        self.bindings: List[int] = []
        self.stream = cuda.Stream()

        # Validate precision input
        if self.precision not in ("fp16", "fp32", "int8"):
            raise ValueError(f"Unsupported precision: {precision}. Use fp16, fp32, or int8.")

        # Build or load existing engine
        if os.path.exists(engine_path):
            self._load_engine()
        else:
            self._build_engine()
            self._save_engine()

    def _build_engine(self) -> None:
        """Build TensorRT engine from ONNX model with specified precision."""
        try:
            builder = trt.Builder(self.logger)
            network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
            parser = trt.OnnxParser(network, self.logger)

            # Load ONNX model
            with open(self.onnx_path, "rb") as f:
                if not parser.parse(f.read()):
                    for i in range(parser.num_errors):
                        print(f"ONNX Parse Error {i}: {parser.get_error(i).desc()}")
                    raise RuntimeError("Failed to parse ONNX model")

            # Configure builder flags
            config = builder.create_builder_config()
            config.max_workspace_size = 1 << 30  # 1GB workspace
            if self.precision == "fp16" and builder.platform_has_fast_fp16:
                config.set_flag(trt.BuilderFlag.FP16)
            elif self.precision == "int8" and builder.platform_has_fast_int8:
                config.set_flag(trt.BuilderFlag.INT8)
                # Note: INT8 requires calibration, omitted here for brevity
            elif self.precision == "fp32":
                pass
            else:
                print(f"Warning: {self.precision} not supported, falling back to FP32")

            # Build engine
            print(f"Building TensorRT engine for {self.onnx_path}...")
            self.engine = builder.build_engine(network, config)
            if not self.engine:
                raise RuntimeError("Failed to build TensorRT engine")
            print("Engine built successfully")

        except Exception as e:
            raise RuntimeError(f"Engine build failed: {str(e)}")

    def _save_engine(self) -> None:
        """Serialize engine to disk for reuse."""
        if not self.engine:
            raise RuntimeError("No engine to save")
        with open(self.engine_path, "wb") as f:
            f.write(self.engine.serialize())
        print(f"Engine saved to {self.engine_path}")

    def _load_engine(self) -> None:
        """Load serialized TensorRT engine from disk."""
        try:
            with open(self.engine_path, "rb") as f:
                runtime = trt.Runtime(self.logger)
                self.engine = runtime.deserialize_cuda_engine(f.read())
            if not self.engine:
                raise RuntimeError("Failed to load engine from disk")
            print(f"Engine loaded from {self.engine_path}")
        except Exception as e:
            raise RuntimeError(f"Engine load failed: {str(e)}")

    def allocate_buffers(self) -> None:
        """Allocate GPU buffers for inputs and outputs."""
        self.inputs = []
        self.outputs = []
        self.bindings = []

        for binding in self.engine:
            shape = self.engine.get_binding_shape(binding)
            size = trt.volume(shape) * self.engine.max_batch_size
            dtype = trt.nptype(self.engine.get_binding_dtype(binding))
            # Allocate host and device buffers
            host_mem = np.empty(size, dtype=dtype)
            device_mem = cuda.mem_alloc(host_mem.nbytes)
            self.bindings.append(int(device_mem))

            if self.engine.binding_is_input(binding):
                self.inputs.append((host_mem, device_mem))
            else:
                self.outputs.append((host_mem, device_mem))

        self.context = self.engine.create_execution_context()

    def infer(self, input_data: np.ndarray) -> np.ndarray:
        """Run inference on input data, return output."""
        if not self.context:
            self.allocate_buffers()

        # Copy input data to device
        np.copyto(self.inputs[0][0], input_data.flatten())
        cuda.memcpy_htod_async(self.inputs[0][1], self.inputs[0][0], self.stream)

        # Run inference
        self.context.execute_async_v2(bindings=self.bindings, stream_handle=self.stream.handle)

        # Copy output back to host
        cuda.memcpy_dtoh_async(self.outputs[0][0], self.outputs[0][1], self.stream)
        self.stream.synchronize()

        return self.outputs[0][0].reshape(self.engine.get_binding_shape(1))

    def benchmark(self, input_shape: tuple, num_runs: int = 10000) -> dict:
        """Benchmark inference latency and throughput."""
        dummy_input = np.random.randn(*input_shape).astype(np.float32)
        latencies = []

        # Warmup
        for _ in range(1000):
            self.infer(dummy_input)

        # Measured runs
        for _ in range(num_runs):
            start = time.perf_counter()
            self.infer(dummy_input)
            end = time.perf_counter()
            latencies.append((end - start) * 1000)  # ms

        return {
            "p50_latency_ms": np.percentile(latencies, 50),
            "p99_latency_ms": np.percentile(latencies, 99),
            "throughput_inf_sec": num_runs / (sum(latencies) / 1000)
        }

if __name__ == "__main__":
    try:
        # Initialize TensorRT inference for ResNet-50
        trt_inf = TensorRTInference(
            onnx_path="resnet50.onnx",
            engine_path="resnet50_fp16.engine",
            precision="fp16"
        )
        trt_inf.allocate_buffers()

        # Run benchmark
        results = trt_inf.benchmark(input_shape=(1, 3, 224, 224), num_runs=10000)
        print(f"TensorRT Benchmark Results: {results}")
    except Exception as e:
        print(f"Fatal error: {str(e)}")
        exit(1)

import onnxruntime as ort
import numpy as np
import time
import os
from typing import List, Dict, Optional

class ONNXRuntimeInference:
    """Production-grade ONNX Runtime inference wrapper with error handling and benchmarking."""

    def __init__(self, onnx_path: str, providers: Optional[List[str]] = None):
        self.onnx_path = onnx_path
        self.providers = providers or ort.get_available_providers()
        self.session: Optional[ort.InferenceSession] = None
        self.input_name: Optional[str] = None
        self.output_name: Optional[str] = None

        # Validate ONNX file exists
        if not os.path.exists(onnx_path):
            raise FileNotFoundError(f"ONNX model not found at {onnx_path}")

        # Initialize session with specified providers
        self._init_session()

    def _init_session(self) -> None:
        """Initialize ONNX Runtime inference session with error handling."""
        try:
            sess_options = ort.SessionOptions()
            sess_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
            sess_options.intra_op_num_threads = 1  # Pin to single thread for consistent benchmarking
            sess_options.log_severity_level = 3  # Suppress info logs

            self.session = ort.InferenceSession(
                path_or_bytes=self.onnx_path,
                sess_options=sess_options,
                providers=self.providers
            )

            # Get input/output names
            self.input_name = self.session.get_inputs()[0].name
            self.output_name = self.session.get_outputs()[0].name
            print(f"ONNX Runtime session initialized with providers: {self.session.get_providers()}")

        except Exception as e:
            raise RuntimeError(f"Failed to initialize ONNX Runtime session: {str(e)}")

    def infer(self, input_data: np.ndarray) -> np.ndarray:
        """Run inference on input data, return output."""
        if not self.session:
            raise RuntimeError("Inference session not initialized")

        try:
            # ONNX Runtime expects input as a dict of name: array
            return self.session.run(
                output_names=[self.output_name],
                input_feed={self.input_name: input_data}
            )[0]
        except Exception as e:
            raise RuntimeError(f"Inference failed: {str(e)}")

    def benchmark(self, input_shape: tuple, num_runs: int = 10000, warmup_runs: int = 1000) -> Dict[str, float]:
        """Benchmark inference latency and throughput."""
        # Validate input shape matches model expectations
        model_input_shape = self.session.get_inputs()[0].shape
        if model_input_shape[1:] != input_shape[1:]:
            raise ValueError(f"Input shape {input_shape} does not match model input shape {model_input_shape}")

        dummy_input = np.random.randn(*input_shape).astype(np.float32)
        latencies = []

        # Warmup runs to stabilize performance
        print(f"Running {warmup_runs} warmup inferences...")
        for _ in range(warmup_runs):
            self.infer(dummy_input)

        # Measured runs
        print(f"Running {num_runs} measured inferences...")
        for _ in range(num_runs):
            start = time.perf_counter()
            self.infer(dummy_input)
            end = time.perf_counter()
            latencies.append((end - start) * 1000)  # Convert to ms

        # Calculate metrics
        return {
            "p50_latency_ms": float(np.percentile(latencies, 50)),
            "p99_latency_ms": float(np.percentile(latencies, 99)),
            "throughput_inf_sec": num_runs / (sum(latencies) / 1000),
            "avg_latency_ms": float(np.mean(latencies))
        }

    def get_model_info(self) -> Dict[str, any]:
        """Return metadata about the loaded ONNX model."""
        if not self.session:
            raise RuntimeError("Session not initialized")
        return {
            "model_path": self.onnx_path,
            "input_names": [inp.name for inp in self.session.get_inputs()],
            "output_names": [out.name for out in self.session.get_outputs()],
            "input_shapes": [inp.shape for inp in self.session.get_inputs()],
            "providers": self.session.get_providers()
        }

if __name__ == "__main__":
    try:
        # Initialize ONNX Runtime for ResNet-50 with CUDA provider
        onnx_inf = ONNXRuntimeInference(
            onnx_path="resnet50.onnx",
            providers=["CUDAExecutionProvider", "CPUExecutionProvider"]
        )

        # Print model info
        print(f"Model Info: {onnx_inf.get_model_info()}")

        # Run benchmark
        results = onnx_inf.benchmark(input_shape=(1, 3, 224, 224), num_runs=10000)
        print(f"ONNX Runtime Benchmark Results: {results}")

    except FileNotFoundError as e:
        print(f"File error: {str(e)}")
        exit(1)
    except RuntimeError as e:
        print(f"Runtime error: {str(e)}")
        exit(1)
    except Exception as e:
        print(f"Unexpected error: {str(e)}")
        exit(1)

import sys
import json
from tensorrt_inference import TensorRTInference
from onnxruntime_inference import ONNXRuntimeInference

def run_comparison_benchmark(
    onnx_path: str,
    trt_engine_path: str,
    input_shape: tuple,
    num_runs: int = 10000,
    precision: str = "fp16"
) -> dict:
    """Run side-by-side benchmark of TensorRT and ONNX Runtime, return structured results."""
    results = {
        "model": onnx_path,
        "input_shape": input_shape,
        "num_runs": num_runs,
        "precision": precision,
        "tensorrt": {},
        "onnxruntime": {}
    }

    # Benchmark TensorRT
    print("\n=== Running TensorRT Benchmark ===")
    try:
        trt_inf = TensorRTInference(
            onnx_path=onnx_path,
            engine_path=trt_engine_path,
            precision=precision
        )
        trt_inf.allocate_buffers()
        trt_results = trt_inf.benchmark(input_shape=input_shape, num_runs=num_runs)
        results["tensorrt"] = trt_results
        print(f"TensorRT Results: {trt_results}")
    except Exception as e:
        print(f"TensorRT benchmark failed: {str(e)}")
        results["tensorrt"]["error"] = str(e)

    # Benchmark ONNX Runtime
    print("\n=== Running ONNX Runtime Benchmark ===")
    try:
        onnx_inf = ONNXRuntimeInference(
            onnx_path=onnx_path,
            providers=["CUDAExecutionProvider", "CPUExecutionProvider"]
        )
        onnx_results = onnx_inf.benchmark(input_shape=input_shape, num_runs=num_runs)
        results["onnxruntime"] = onnx_results
        print(f"ONNX Runtime Results: {onnx_results}")
    except Exception as e:
        print(f"ONNX Runtime benchmark failed: {str(e)}")
        results["onnxruntime"]["error"] = str(e)

    # Calculate relative performance
    if "p99_latency_ms" in results["tensorrt"] and "p99_latency_ms" in results["onnxruntime"]:
        trt_lat = results["tensorrt"]["p99_latency_ms"]
        onnx_lat = results["onnxruntime"]["p99_latency_ms"]
        results["relative_performance"] = {
            "trt_vs_onnx_latency_ratio": trt_lat / onnx_lat,
            "onnx_vs_trt_throughput_ratio": results["onnxruntime"].get("throughput_inf_sec", 0) / results["tensorrt"].get("throughput_inf_sec", 1)
        }

    return results

def save_results(results: dict, output_path: str = "benchmark_results.json") -> None:
    """Save benchmark results to JSON file."""
    try:
        with open(output_path, "w") as f:
            json.dump(results, f, indent=2)
        print(f"\nResults saved to {output_path}")
    except Exception as e:
        print(f"Failed to save results: {str(e)}")

def print_markdown_table(results: dict) -> None:
    """Print benchmark results as a Markdown comparison table."""
    print("\n=== Comparison Table ===")
    print("| Metric                | TensorRT 10.1 | ONNX Runtime 1.17 | Winner       |")
    print("|-----------------------|---------------|-------------------|--------------|")

    metrics = ["p50_latency_ms", "p99_latency_ms", "throughput_inf_sec"]
    labels = ["p50 Latency (ms)", "p99 Latency (ms)", "Throughput (inf/sec)"]

    for metric, label in zip(metrics, labels):
        trt_val = results["tensorrt"].get(metric, "N/A")
        onnx_val = results["onnxruntime"].get(metric, "N/A")

        if trt_val != "N/A" and onnx_val != "N/A":
            if metric == "throughput_inf_sec":
                winner = "TensorRT" if trt_val > onnx_val else "ONNX Runtime"
            else:
                winner = "TensorRT" if trt_val < onnx_val else "ONNX Runtime"
        else:
            winner = "N/A"

        print(f"| {label:<21} | {trt_val:<13} | {onnx_val:<17} | {winner:<12} |")

if __name__ == "__main__":
    # Validate command line arguments
    if len(sys.argv) < 2:
        print("Usage: python compare_benchmarks.py  [input_shape] [num_runs]")
        print("Example: python compare_benchmarks.py resnet50.onnx 1,3,224,224 10000")
        exit(1)

    onnx_path = sys.argv[1]
    trt_engine_path = f"{onnx_path}.{sys.argv[2] if len(sys.argv) > 2 else 'fp16'}.engine"

    # Parse input shape
    if len(sys.argv) > 2:
        input_shape = tuple(map(int, sys.argv[2].split(",")))
    else:
        input_shape = (1, 3, 224, 224)

    # Parse num runs
    num_runs = int(sys.argv[3]) if len(sys.argv) > 3 else 10000

    # Run benchmark
    results = run_comparison_benchmark(
        onnx_path=onnx_path,
        trt_engine_path=trt_engine_path,
        input_shape=input_shape,
        num_runs=num_runs
    )

    # Print table
    print_markdown_table(results)

    # Save results
    save_results(results)

# Short snippet to check if your custom layer is supported by TensorRT
import tensorrt as trt
logger = trt.Logger(trt.Logger.WARNING)
builder = trt.Builder(logger)
network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
parser = trt.OnnxParser(network, logger)
with open("your_custom_model.onnx", "rb") as f:
    if not parser.parse(f.read()):
        print("Unsupported layers found:")
        for i in range(parser.num_errors):
            print(f"  - {parser.get_error(i).desc()}")

# Check available execution providers for ONNX Runtime
import onnxruntime as ort
print("Available ONNX Runtime Providers:")
for provider in ort.get_available_providers():
    print(f"  - {provider}")
# Output on NVIDIA system: CUDAExecutionProvider, CPUExecutionProvider
# Output on AMD system: ROCMExecutionProvider, CPUExecutionProvider

# GitHub Actions step to pre-build TensorRT engines
- name: Build TensorRT Engines
  run: |
    python tensorrt_inference.py --build-only --onnx-model resnet50.onnx --precision fp16 --engine-path engines/resnet50_fp16.engine
    python tensorrt_inference.py --build-only --onnx-model resnet50.onnx --precision int8 --engine-path engines/resnet50_int8.engine
- name: Upload Engines to S3
  uses: aws-actions/configure-aws-credentials@v4
  with:
    aws-access-key-id: ${{ secrets.AWS_ACCESS_KEY_ID }}
    aws-secret-access-key: ${{ secrets.AWS_SECRET_ACCESS_KEY }}
    aws-region: us-east-1
- run: aws s3 sync engines/ s3://my-model-registry/tensorrt-engines/