{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.SPDX-License-Identifier: Apache-2.0\n",
    "\n",
    "Licensed under the Apache License, Version 2.0 (the \"License\"); you may not use\n",
    "\n",
    "this file except in compliance with the License. You may obtain a copy of the License at\n",
    "\n",
    "\n",
    "\n",
    "http://www.apache.org/licenses/LICENSE-2.0\n",
    "\n",
    "\n",
    "\n",
    "Unless required by applicable law or agreed to in writing, software\n",
    "\n",
    "distributed under the License is distributed on an \"AS IS\" BASIS,\n",
    "\n",
    "WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
    "\n",
    "See the License for the specific language governing permissions and\n",
    "\n",
    "limitations under the License.\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Getting Started with TensorRT: Accelerate Your Deep Learning Inference\n",
    "\n",
    "Welcome to your first TensorRT tutorial! In this notebook, you'll learn how to:\n",
    "1. Load a pre-trained EfficientNet model in ONNX format\n",
    "2. Convert it to a TensorRT engine for faster inference\n",
    "3. Run inference and see the speedup firsthand\n",
    "4. Make predictions on real images\n",
    "\n",
    "## Understanding ONNX: The Universal Model Format\n",
    "\n",
    "ONNX (Open Neural Network Exchange) is a standard format for representing deep learning models. Think of it as a universal language that different deep learning frameworks can understand. Here's why it's important:\n",
    "\n",
    "- **Framework Independence**: Models trained in PyTorch, TensorFlow, or other frameworks can be exported to ONNX\n",
    "- **Interoperability**: ONNX models can be imported into various inference engines and frameworks\n",
    "- **Production Ready**: ONNX is widely used in production environments for model deployment\n",
    "\n",
    "### The ONNX to TensorRT Workflow\n",
    "\n",
    "TensorRT is NVIDIA's deep learning inference optimizer that can import models from ONNX. This makes it a powerful tool in your deployment pipeline:\n",
    "\n",
    "```\n",
    "Your Framework (PyTorch/TF/etc.) → ONNX → TensorRT ===> Optimized Inference\n",
    "```\n",
    "\n",
    "This workflow is particularly powerful because:\n",
    "1. You can train your model in any framework you prefer\n",
    "2. Export it to ONNX (a one-time conversion)\n",
    "3. Use TensorRT to optimize it for NVIDIA GPUs\n",
    "4. Get significant speedup in production\n",
    "\n",
    "## Prerequisites\n",
    "\n",
    "Before we start, make sure you have:\n",
    "- NVIDIA GPU with CUDA support\n",
    "- Python 3.10+ installed\n",
    "- Basic understanding of deep learning and inference\n",
    "\n",
    "Let's begin by installing and importing the required packages:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install tensorrt cuda-python pillow onnxruntime\n",
    "import tensorrt as trt\n",
    "from cuda.bindings import runtime as cudart\n",
    "from PIL import Image\n",
    "import numpy as np\n",
    "from pathlib import Path\n",
    "import time\n",
    "from typing import Optional, Union, Tuple"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "root = Path.cwd()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# define a function to download files\n",
    "\n",
    "import requests\n",
    "from requests.adapters import HTTPAdapter\n",
    "from urllib3.util.retry import Retry\n",
    "\n",
    "def download_file(url: str, output_path: Union[str, Path]):\n",
    "    \"\"\"Download a file with retry mechanism.\"\"\"\n",
    "    session = requests.Session()\n",
    "    retry = Retry(total=10, backoff_factor=1)\n",
    "    adapter = HTTPAdapter(max_retries=retry)\n",
    "    session.mount('http://', adapter)\n",
    "    session.mount('https://', adapter)\n",
    "\n",
    "    response = session.get(url, verify=False, timeout=30)\n",
    "    output_path.parent.mkdir(parents=True, exist_ok=True)\n",
    "    with open(output_path, 'wb') as f:\n",
    "        f.write(response.content)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 1: Download a Pre-trained Model\n",
    "\n",
    "We'll use EfficientNet-B0, a popular and efficient image classification model, as an example for this sample. \n",
    "\n",
    "### Understanding ONNX Model Structure\n",
    "\n",
    "An ONNX model contains:\n",
    "- Model architecture (layers, connections)\n",
    "- Weights and biases\n",
    "- Input/output specifications\n",
    "- Metadata about the model\n",
    "just like any other model representations. \n",
    "\n",
    "This standardized format makes it easy to move models between different frameworks and inference engines."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "download_file(\"https://github.com/onnx/models/raw/refs/heads/main/Computer_Vision/efficientnet_b0_Opset17_timm/efficientnet_b0_Opset17.onnx\", root / \"efficientnet-b0.onnx\")\n",
    "assert (root / \"efficientnet-b0.onnx\").exists(), \"Model file not found. Please check if the download was successful.\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 2: Convert ONNX to TensorRT Engine\n",
    "\n",
    "This is where the magic happens! We'll convert our ONNX model into a TensorRT engine. The engine is optimized for your specific GPU and will run much faster than the original model.\n",
    "\n",
    "### The Conversion Process\n",
    "\n",
    "1. **Load ONNX Model**: TensorRT reads the ONNX file and understands the model structure\n",
    "2. **Optimize**: TensorRT performs several optimizations:\n",
    "   - Layer fusion\n",
    "   - Memory optimization\n",
    "   - Precision calibration\n",
    "3. **Generate Engine**: Creates a highly optimized inference engine\n",
    "\n",
    "The resulting engine is specific to your GPU and will run much faster than the original ONNX model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "logger = trt.Logger(trt.Logger.WARNING)\n",
    "builder = trt.Builder(logger)\n",
    "network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.STRONGLY_TYPED))\n",
    "\n",
    "# Bind the TensorRT network to the parser so that the parser can update the network later accordingly\n",
    "parser = trt.OnnxParser(network, logger)\n",
    "\n",
    "onnx_path = root / \"efficientnet-b0.onnx\"\n",
    "print(f'Parsing ONNX model at {onnx_path}...')\n",
    "with open(onnx_path, \"rb\") as model:\n",
    "    parser.parse(model.read())\n",
    "print('Parsing ONNX model... done')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now that we have the TensorRT `INetworkDefinition`, we can start building the engine"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "config = builder.create_builder_config()\n",
    "\n",
    "# TensorRT needs memory for layer operations and intermediate activations during inference\n",
    "# Setting a memory limit helps control resource usage and prevents out-of-memory errors\n",
    "config.set_memory_pool_limit(\n",
    "        trt.MemoryPoolType.WORKSPACE, 1 << 30\n",
    ") # 1GB\n",
    "\n",
    "print('Starting to build engine. This might take several minutes depending on the hardware...')\n",
    "engine = builder.build_serialized_network(network, config)\n",
    "assert engine is not None, 'Engine build failed'\n",
    "\n",
    "engine_path = root / \"efficientnet-b0.plan\"\n",
    "with open(engine_path, 'wb') as f:\n",
    "    f.write(engine)\n",
    "\n",
    "print(\"TensorRT engine created successfully!\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Optional: Using Editable Timing Cache\n",
    "\n",
    "TensorRT engines may vary between builds because kernel selection is based on runtime performance measurements. The hardware state (GPU utilization, temperature, system load) affects which kernels are chosen since kernels might outperform each other under different scenarios. \n",
    "\n",
    "To ensure consistent builds, TensorRT provides an editable timing cache that:\n",
    "- Stores intermediate optimization results\n",
    "- Enables deterministic engine builds\n",
    "- Speeds up subsequent builds since they don't need to measure kernel execution time for each op again"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def build_engine_with_cache(onnx_path: Union[str, Path], timing_cache: Optional[trt.ITimingCache]):\n",
    "    builder = trt.Builder(logger)\n",
    "    network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.STRONGLY_TYPED))\n",
    "    parser = trt.OnnxParser(network, logger)\n",
    "    with open(onnx_path, 'rb') as model:\n",
    "        parser.parse(model.read())\n",
    "    config = builder.create_builder_config()\n",
    "    config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 1 << 30)\n",
    "\n",
    "    # Enable editable timing cache\n",
    "    config.set_flag(trt.BuilderFlag.EDITABLE_TIMING_CACHE)\n",
    "\n",
    "    # Create timing cache if not provided\n",
    "    if not timing_cache:\n",
    "        timing_cache = config.create_timing_cache(bytes())\n",
    "    config.set_timing_cache(timing_cache, True)\n",
    "\n",
    "    # Build engine\n",
    "    print('Start building engine...')\n",
    "    tik = time.time()\n",
    "    engine = builder.build_serialized_network(network, config)\n",
    "    tok = time.time()\n",
    "\n",
    "    print(f'Engine build cost {tok - tik}ms')\n",
    "    return engine, timing_cache\n",
    "\n",
    "# First build (creates cache)\n",
    "engine1, timing_cache = build_engine_with_cache(onnx_path, None)\n",
    "print(\"First build completed with cache creation\")\n",
    "\n",
    "# Second build (uses cache)\n",
    "engine2, timing_cache = build_engine_with_cache(onnx_path, timing_cache)\n",
    "print(\"Second build completed with cache creation\")\n",
    "\n",
    "is_identical = np.array_equal(\n",
    "    np.frombuffer(engine1, dtype=np.uint8),\n",
    "    np.frombuffer(engine2, dtype=np.uint8))\n",
    "print(f'Is engine identical: {is_identical}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 3: Run Inference and Compare Performance\n",
    "\n",
    "Now let's see the real power of TensorRT! We'll:\n",
    "1. Run inference with both ONNX and TensorRT\n",
    "2. Compare their performance\n",
    "3. See the speedup TensorRT provides\n",
    "\n",
    "### Understanding the Performance Difference\n",
    "\n",
    "The speedup comes from several optimizations:\n",
    "- Layer fusion: Combining multiple operations into one\n",
    "- Memory optimization: Better memory access patterns\n",
    "- Precision optimization: Using optimal precision for each layer\n",
    "- CUDA optimization: Direct GPU execution without framework overhead"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_and_preprocess_image(image_path: Union[str, Path], input_size: Tuple[int, int] = (224, 224)):\n",
    "    img = Image.open(image_path)\n",
    "    img = img.resize(input_size)\n",
    "    img = np.array(img).astype(np.float32)\n",
    "    img = img / 255.0  # Normalize from [0, 255] to [0, 1]\n",
    "    img = np.transpose(img, (2, 0, 1))  # HWC to CHW\n",
    "    img = np.expand_dims(img, axis=0)  # Add batch dimension\n",
    "    return img\n",
    "\n",
    "def check_cuda_error(error):\n",
    "    if isinstance(error, tuple):\n",
    "        error = error[0]\n",
    "    if error != cudart.cudaError_t.cudaSuccess:\n",
    "        error_name = cudart.cudaGetErrorName(error)[1]\n",
    "        error_string = cudart.cudaGetErrorString(error)[1]\n",
    "        raise RuntimeError(f\"CUDA Error: {error_name} ({error_string})\")\n",
    "\n",
    "def run_inference_trt(engine: trt.ICudaEngine, input_data: np.ndarray):\n",
    "    # Create execution context - this stores the device memory allocations\n",
    "    # and bindings needed for inference\n",
    "    context = engine.create_execution_context()\n",
    "\n",
    "    # Initialize lists to store input/output information and GPU memory allocations\n",
    "    inputs = []\n",
    "    outputs = []\n",
    "    allocations = []\n",
    "\n",
    "    # Iterate through all input/output tensors to set up memory and bindings\n",
    "    for i in range(engine.num_io_tensors):\n",
    "        name = engine.get_tensor_name(i)\n",
    "        # Check if this tensor is an input or output\n",
    "        is_input = engine.get_tensor_mode(name) == trt.TensorIOMode.INPUT\n",
    "        # Get tensor datatype and shape information\n",
    "        dtype = engine.get_tensor_dtype(name)\n",
    "        shape = engine.get_tensor_shape(name)\n",
    "\n",
    "        # Calculate required memory size for this tensor\n",
    "        size = np.dtype(trt.nptype(dtype)).itemsize\n",
    "        for s in shape:\n",
    "            size *= s\n",
    "\n",
    "        # Allocate GPU memory for this tensor\n",
    "        err, allocation = cudart.cudaMalloc(size)\n",
    "        check_cuda_error(err)\n",
    "\n",
    "        # Store tensor information in a dictionary for easy access\n",
    "        binding = {\n",
    "            \"index\": i,\n",
    "            \"name\": name,\n",
    "            \"dtype\": np.dtype(trt.nptype(dtype)),\n",
    "            \"shape\": list(shape),\n",
    "            \"allocation\": allocation,\n",
    "            \"size\": size,\n",
    "        }\n",
    "\n",
    "        # Keep track of all allocations and sort tensors into inputs/outputs\n",
    "        allocations.append(allocation)\n",
    "        if is_input:\n",
    "            inputs.append(binding)\n",
    "        else:\n",
    "            outputs.append(binding)\n",
    "\n",
    "    # Ensure input data is contiguous in memory for efficient GPU transfer\n",
    "    input_data = np.ascontiguousarray(input_data)\n",
    "\n",
    "    # Copy input data from host (CPU) to device (GPU)\n",
    "    err = cudart.cudaMemcpy(\n",
    "        inputs[0][\"allocation\"],\n",
    "        input_data.ctypes.data,\n",
    "        inputs[0][\"size\"],\n",
    "        cudart.cudaMemcpyKind.cudaMemcpyHostToDevice,\n",
    "    )\n",
    "    check_cuda_error(err)\n",
    "\n",
    "    # Set tensor addresses for all tensors\n",
    "    for i in range(engine.num_io_tensors):\n",
    "        context.set_tensor_address(engine.get_tensor_name(i), allocations[i])\n",
    "\n",
    "    # Create a CUDA stream for asynchronous execution\n",
    "    err, stream = cudart.cudaStreamCreate()\n",
    "    check_cuda_error(err)\n",
    "\n",
    "    # Run inference using the TensorRT engine\n",
    "    context.execute_async_v3(stream_handle=stream)\n",
    "    err = cudart.cudaStreamSynchronize(stream)\n",
    "    check_cuda_error(err)\n",
    "\n",
    "    # Prepare numpy array for output and copy results from GPU to CPU\n",
    "    output_shape = outputs[0][\"shape\"]\n",
    "    output = np.empty(output_shape, dtype=outputs[0][\"dtype\"])\n",
    "\n",
    "    err = cudart.cudaMemcpy(\n",
    "        output.ctypes.data,\n",
    "        outputs[0][\"allocation\"],\n",
    "        outputs[0][\"size\"],\n",
    "        cudart.cudaMemcpyKind.cudaMemcpyDeviceToHost,\n",
    "    )\n",
    "    check_cuda_error(err)\n",
    "\n",
    "    # Free all GPU memory allocations\n",
    "    for allocation in allocations:\n",
    "        err = cudart.cudaFree(allocation)\n",
    "        check_cuda_error(err)\n",
    "\n",
    "    # Destroy the CUDA stream\n",
    "    err = cudart.cudaStreamDestroy(stream)\n",
    "    check_cuda_error(err)\n",
    "\n",
    "    return output\n",
    "\n",
    "import onnxruntime as ort\n",
    "def run_inference_onnx(session, input_data: np.ndarray):\n",
    "    output = session.run(None, {'x': input_data})[0]\n",
    "    return output"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Let's Compare Performance!\n",
    "\n",
    "We'll run both models multiple times to get an accurate comparison of their performance. This will show you the baseline speedup that TensorRT provides. \n",
    "\n",
    "Refer to https://docs.nvidia.com/deeplearning/tensorrt/latest/index.html for more information about how to further optimize your engine"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create a sample input\n",
    "sample_input = np.random.randn(1, 3, 224, 224).astype(np.float32)\n",
    "\n",
    "# Benchmark ONNX Runtime\n",
    "session = ort.InferenceSession(onnx_path)\n",
    "onnx_times = []\n",
    "for _ in range(100):\n",
    "    start_time = time.time()\n",
    "    _ = run_inference_onnx(session, sample_input)\n",
    "    onnx_times.append(time.time() - start_time)\n",
    "\n",
    "# Benchmark TensorRT\n",
    "with open(engine_path, \"rb\") as f, trt.Runtime(logger) as runtime:\n",
    "    engine = runtime.deserialize_cuda_engine(f.read())\n",
    "trt_times = []\n",
    "for _ in range(100):\n",
    "    start_time = time.time()\n",
    "    _ = run_inference_trt(engine, sample_input)\n",
    "    trt_times.append(time.time() - start_time)\n",
    "\n",
    "print(f\"ONNX Runtime Average Time: {np.mean(onnx_times)*1000:.2f} ms\")\n",
    "print(f\"TensorRT Average Time: {np.mean(trt_times)*1000:.2f} ms\")\n",
    "print(f\"Speedup: {np.mean(onnx_times)/np.mean(trt_times):.2f}x\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 4: Run Inference on a Real Image\n",
    "\n",
    "Now let's try our optimized model on a real image! We'll:\n",
    "1. Download a sample image\n",
    "2. Load the ImageNet class labels\n",
    "3. Make predictions and show the results\n",
    "\n",
    "This will demonstrate how the optimized TensorRT engine performs in a real-world scenario."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Download a sample image\n",
    "download_file(\"https://raw.githubusercontent.com/pytorch/hub/master/images/dog.jpg\", root / \"test_image.jpg\")\n",
    "\n",
    "from PIL import Image\n",
    "from IPython.display import display\n",
    "\n",
    "# Open and display the image\n",
    "img = Image.open(root/\"test_image.jpg\")\n",
    "display(img)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_imagenet_labels():\n",
    "    # Download ImageNet labels if not exists\n",
    "    if not (root / \"imagenet_classes.txt\").is_file():\n",
    "        download_file(\"https://raw.githubusercontent.com/pytorch/hub/master/imagenet_classes.txt\", root / \"imagenet_classes.txt\")\n",
    "    # Read the labels\n",
    "    with open(root / \"imagenet_classes.txt\") as f:\n",
    "        categories = [s.strip() for s in f.readlines()]\n",
    "    return categories\n",
    "\n",
    "# Load ImageNet labels\n",
    "categories = load_imagenet_labels()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load and preprocess a test image\n",
    "test_image_path = root / \"test_image.jpg\"\n",
    "input_data = load_and_preprocess_image(test_image_path)\n",
    "\n",
    "# Run inference\n",
    "output = run_inference_trt(engine, input_data)\n",
    "\n",
    "# Get top 5 predictions\n",
    "top5_idx = np.argsort(output[0])[-5:][::-1]\n",
    "print(\"Top 5 predictions:\")\n",
    "for idx in top5_idx:\n",
    "    print(f\"{categories[idx]}: {output[0][idx]:.2f}%\")\n",
    "assert categories[top5_idx[0]] == \"Samoyed\", 'Incorrect prediction'\n",
    "print('Correctly recognized!')\n",
    "print('Notebook executed successfully')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Congratulations! 🎉\n",
    "\n",
    "### You've successfully:\n",
    "1. Loaded a pre-trained EfficientNet model in ONNX format\n",
    "2. Converted it to a TensorRT engine\n",
    "3. Achieved significant speedup in inference\n",
    "4. Made predictions on real images\n",
    "5. Learned how to use timing cache to speed up engine building and ensure engine build determinism. \n",
    "\n",
    "### What's Next?\n",
    "\n",
    "Now that you understand the ONNX to TensorRT workflow, you can:\n",
    "- Export your own models from PyTorch/TensorFlow to ONNX\n",
    "- Try different optimization settings in TensorRT\n",
    "- Apply this workflow to your production models and get instant performance boost with NVIDIA GPUs!\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
