wehub/nvidia--tensorrt
1
0
Fork 0
Code Issues Pull Requests Actions 26 Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution
Docker Image CI / build-ubuntu2004 (push) Waiting to run
Details

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 13:36:55 +08:00
commit c8a779b1bb
1887 changed files with 3245738 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
samples/python/refactored/1_run_onnx_with_tensorrt/README.md
+66
View File
@@ -0,0 +1,66 @@
# Run ONNX with TensorRT
This sample demonstrates:
- Converting a pre-trained [EfficientNet](https://arxiv.org/abs/1905.11946)-B0 ONNX model to a `TensorRT` engine
- Performing inference with `TensorRT` using Python APIs
- Comparing inference performance between `ONNX Runtime` and `TensorRT`
- Proper memory management and resource cleanup in both Python implementations
## Key features demonstrated:
- `TensorRT`'s ONNX parser + ONNX model -> `TensorRT` engine
- Engine building and serialization
- Input/output tensor handling
- Performance profiling
- Editable timing cache for deterministic engine builds
- Memory pool optimization with workspace configuration
## Implementation Details
### Memory Management
- Configures workspace memory pool for running under limited hardware
### Engine Building
- Supports editable timing cache for deterministic builds
- Serialization and deserialization of TensorRT engines
### Inference Pipeline
- Efficient image preprocessing with `PIL` and `NumPy`
- Supports batch inference
- Implements proper error handling and resource cleanup
- Provides performance comparison between `ONNX Runtime` and `TensorRT`
- Performs inference on a real-world image
## CLI Tools
Users can run their onnx model and generate the engine with similar functionality using `trtexec`:
```bash
# Basic conversion with performance profiling
trtexec --onnx=efficientnet-b0.onnx \
--saveEngine=efficientnet-b0_trtexec.plan \
--dumpProfile \
--iterations=100 \
--avgRuns=100 \
--workspace=1024 \
--batch=1
```
Key options explained:
- `--onnx`: Input ONNX model
- `--saveEngine`: Output TensorRT engine
- `--dumpProfile`: Performance profiling
- `--iterations`: Number of inference iterations
- `--avgRuns`: Number of runs to average for timing
- `--workspace`: Workspace size in MB (1024MB = 1GB)
- `--batch`: Batch size for inference
## Additional Resources
- [TensorRT Documentation](https://docs.nvidia.com/deeplearning/tensorrt/latest/index.html)
- [ONNX Documentation](https://onnx.ai/)
- [EfficientNet Paper](https://arxiv.org/abs/1905.11946)
# Changelog
August 2025
Removed support for Python versions < 3.10.
samples/python/refactored/1_run_onnx_with_tensorrt/main.ipynb
+589
View File
@@ -0,0 +1,589 @@
{
"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
}
samples/python/refactored/2_construct_network_with_layer_apis/README.md
+21
View File
@@ -0,0 +1,21 @@
# Sample 2: Constructing a Network with TensorRT Layer APIs
This sample demonstrates how to build a TensorRT network definition from scratch using the TensorRT Layer APIs, focusing on constructing a recurrent neural network (LSTM) and utilizing advanced builder features.
## Description
This sample constructs a simple, single-layer Long Short-Term Memory (LSTM) network using the TensorRT Layer APIs. The primary goal is to illustrate how to:
1. Define individual network layers and their connections programmatically using Python (**TensorRT Layer API**). This includes layers like constants, matrix multiply, element-wise operations, activations, and slicing.
2. Implement recurrent logic by building an LSTM cell and using TensorRT's `add_loop` construct to create a recurrent LSTM layer (**Recurrent Network Construction**).
3. Monitor the potentially lengthy engine build process by implementing `IProgressMonitor` for real-time feedback (**Build Progress Monitoring**).
4. Configure the builder for engine portability using `BuilderFlag.VERSION_COMPATIBLE` to create more portable engines (**Version-Compatible Engines**).
5. Run inference and verify the custom network's correctness by utilizing Polygraphy's `TrtRunner` for simplified engine loading/execution (**Inference with Polygraphy**) and comparing the TensorRT engine's output against a reference NumPy implementation (**NumPy Verification**).
## Additional Resources
* [tensorrtx repo](https://github.com/wang-xinyu/tensorrtx): Offers real-world examples of constructing complex networks using the TensorRT Layer APIs.
# Changelog
August 2025
Removed support for Python versions < 3.10.
samples/python/refactored/2_construct_network_with_layer_apis/main.ipynb
+694
View File
@@ -0,0 +1,694 @@
{
"cells": [
{
"cell_type": "markdown",
"metadata": {},
"source": [
"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."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# 2. Constructing a Network with TensorRT Layer APIs\n",
"\n",
"In this notebook, you'll learn how to move beyond pre-built model formats and directly construct a neural network using TensorRT's versatile Layer APIs. This approach offers fine-grained control over your network architecture and optimizations.\n",
"\n",
"Specifically, we will cover:\n",
"\n",
"1. **Building a Recurrent Network (LSTM) from Scratch:** Understand how to define each layer of a Long Short-Term Memory (LSTM) cell and then use these components to construct an entire recurrent LSTM layer. This involves using various Layer API functionalities like `add_constant`, `add_matrix_multiply`, `add_elementwise`, `add_slice`, and `add_activation`.\n",
"2. **Implementing Loops for Recurrence:** Utilize TensorRT's `add_loop` functionality to efficiently handle the recurrent nature of the LSTM, processing an input sequence step-by-step.\n",
"3. **Monitoring Build Progress:** Implement an `IProgressMonitor` to track the engine creation process in real-time, providing visibility into potentially long build times.\n",
"4. **Creating Version-Compatible Engines:** Learn to save TensorRT engines with the `BuilderFlag.VERSION_COMPATIBLE` flag, enhancing their portability across different TensorRT patch versions and compatible hardware.\n",
"\n",
"This sample uses a small, single-layer LSTM to keep the focus on these core TensorRT API features. We'll also verify its output against an equivalent NumPy implementation to ensure correctness."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Introduction\n",
"\n",
"While importing models via ONNX offers convenience, constructing networks directly with TensorRT APIs provides fine-grained control over the network definition. The **[TensorRT Layer API](https://docs.nvidia.com/deeplearning/tensorrt/latest/python_api/infer/Graph/Layers.html)** enables users to define each layer explicitly, offering flexibility and optimization opportunities.\n",
"\n",
"To facilitate understanding and verification, this demonstration employs small tensors, allowing for direct comparison with an equivalent NumPy implementation.\n",
"\n",
"> **Note: This sample assumes familiarity with the basic concepts of Long Short-Term Memory (LSTM) networks. If you're new to LSTMs, you might find it helpful to review their structure and operation before proceeding.**"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Step 0: Prerequisites"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"%pip install numpy tensorrt polygraphy --extra-index-url https://pypi.ngc.nvidia.com\n",
"import tensorrt as trt\n",
"import numpy as np\n",
"from typing import Tuple"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"To simplify, network parameters and weight initializations use small, illustrative values and dimensions."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# === Network Parameters & Weights Initialization ===\n",
"batch_size = 1\n",
"seq_len = 5 # Length of the sequence\n",
"input_size = 1 # Dimension of input vector at each time step\n",
"hidden_size = 2 # Dimension of hidden state and cell state\n",
"num_units = 1\n",
"\n",
"# --- Create Fixed Dummy Weights and Biases (NumPy arrays with dummy values) ---\n",
"# These will be used by both the TensorRT build and the NumPy verification\n",
"w_val, u_val, b_val = 0.01, 0.05, 0.3\n",
"initial_h_val = 0.1\n",
"initial_c_val = 0.2\n",
"\n",
"# Define shapes\n",
"w_shape = (input_size, 4 * hidden_size) # e.g., [1, 8] for layer 0\n",
"u_shape = (hidden_size, 4 * hidden_size) # e.g., [2, 8]\n",
"b_shape = (4 * hidden_size,) # e.g., [8]\n",
"initial_h_shape = (batch_size, hidden_size)\n",
"initial_c_shape = (batch_size, hidden_size)\n",
"\n",
"# Create NumPy arrays\n",
"np_weight_W = np.full(w_shape, w_val, dtype=np.float32)\n",
"np_weight_U = np.full(u_shape, u_val, dtype=np.float32)\n",
"np_bias = np.full(b_shape, b_val, dtype=np.float32)\n",
"np_initial_h = np.full(initial_h_shape, initial_h_val, dtype=np.float32)\n",
"np_initial_c = np.full(initial_c_shape, initial_c_val, dtype=np.float32)\n",
"\n",
"# Create inputs for the network\n",
"np_inputs = np.ones((seq_len, batch_size, input_size), dtype=np.float32)\n",
"\n",
"print(\"NumPy Weights Initialized:\")\n",
"print(f\" W shape : {np_weight_W.shape}\")\n",
"print(f\" U shape : {np_weight_U.shape}\")\n",
"print(f\" Bias shape : {np_bias.shape}\")\n",
"print(f\" Initial H shape : {np_initial_h.shape}\")\n",
"print(f\" Initial C shape : {np_initial_c.shape}\")\n"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Step 1: Defining LSTM Operations with the Layer API\n",
"\n",
"This step involves defining the LSTM operations by adding layers to the TensorRT `INetworkDefinition`.\n",
"\n",
"### Typical Usage Pattern for TensorRT Layer APIs\n",
"\n",
"When adding layers to a TensorRT network using the Layer API, the common pattern is:\n",
"\n",
"1. **Add the layer:** Use a `network.add_*` method (e.g., `network.add_matrix_multiply`) to add the desired layer. This method takes input tensors and layer-specific parameters, returning an `ILayer` object representing the newly added layer.\n",
"2. **Configure the layer:** Access the returned `ILayer` object to configure its properties. This is optional but useful for naming layer's name and output tensors name for easier debugging and more helpful logs. \n",
"\n",
"```python\n",
"# Example: Adding and configuring a generic layer\n",
"\n",
"# 1. Add the layer (replace with a specific layer like add_matrix_multiply)\n",
"layer = network.add_some_layer(input_tensor, ...)\n",
"\n",
"# 2. Configure the layer (optional)\n",
"output_tensor = layer.get_output(0)\n",
"output_tensor.name = 'my_layer_output' # Name the output\n",
"# ... other configurations ...\n",
"```\n",
"\n",
"For a comprehensive list of available layer types and their specific methods and properties, consult the official [TensorRT Layer API documentation](https://docs.nvidia.com/deeplearning/tensorrt/latest/_static/python-api/infer/Graph/Layers.html)."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"TRT_LOGGER = trt.Logger(trt.Logger.INFO)\n",
"\n",
"def add_lstm_unit(network: trt.INetworkDefinition,\n",
" input_x: trt.ITensor, # Shape: [batch_size, input_size]\n",
" prev_h: trt.ITensor, # Shape: [batch_size, hidden_size]\n",
" prev_c: trt.ITensor, # Shape: [batch_size, hidden_size]\n",
" W: np.ndarray, # Shape: [input_size, 4 * hidden_size]\n",
" U: np.ndarray, # Shape: [hidden_size, 4 * hidden_size]\n",
" bias: np.ndarray, # Shape: [4 * hidden_size]\n",
" hidden_size: int,\n",
" input_size: int\n",
" ) -> Tuple[trt.ITensor, trt.ITensor]:\n",
" \"\"\"\n",
" Adds the computations for a single LSTM time step.\n",
" Assumes input tensors have a leading batch dimension.\n",
" \"\"\"\n",
" batch_size = input_x.shape[0] # Get batch size from input\n",
"\n",
" # Create constant layers for weights and biases\n",
" W_layer = network.add_constant(W.shape, trt.Weights(W))\n",
" W_layer.get_output(0).name = \"W_const\"\n",
" U_layer = network.add_constant(U.shape, U)\n",
" U_layer.get_output(0).name = \"U_const\"\n",
" # Reshape bias for broadcasting: [4*hidden] -> [1, 4*hidden]\n",
" bias_reshaped_np = np.expand_dims(bias.copy(), axis=0)\n",
" bias_layer = network.add_constant(bias_reshaped_np.shape, bias_reshaped_np)\n",
" bias_layer.get_output(0).name = \"Bias_const\"\n",
"\n",
"\n",
" # Linear transformations: Wx = input_x * W ; Uh = prev_h * U\n",
" # Wx = [batch, input] * [input, 4*hidden] = [batch, 4*hidden]\n",
" mm_wx = network.add_matrix_multiply(input_x, trt.MatrixOperation.NONE,\n",
" W_layer.get_output(0), trt.MatrixOperation.NONE)\n",
" mm_wx.get_output(0).name = \"Wx\"\n",
"\n",
" # Uh = [batch, hidden] * [hidden, 4*hidden] = [batch, 4*hidden]\n",
" mm_uh = network.add_matrix_multiply(prev_h, trt.MatrixOperation.NONE,\n",
" U_layer.get_output(0), trt.MatrixOperation.NONE)\n",
" mm_uh.get_output(0).name = \"Uh\"\n",
"\n",
"\n",
" # Combined gates = Wx + Uh + Bias\n",
" gates_wx_uh = network.add_elementwise(mm_wx.get_output(0), mm_uh.get_output(0),\n",
" trt.ElementWiseOperation.SUM)\n",
" gates_wx_uh.get_output(0).name = \"Wx_plus_Uh\"\n",
"\n",
" gates = network.add_elementwise(gates_wx_uh.get_output(0), bias_layer.get_output(0),\n",
" trt.ElementWiseOperation.SUM)\n",
"\n",
" gates_output = gates.get_output(0) # Shape [batch, 4*hidden]\n",
" gates_output.name = \"Gates_Combined\"\n",
"\n",
" # Split the combined gates tensor [batch, 4*hidden] -> four [batch, hidden] gate tensors (Input, Forget, Candidate, Output)\n",
" def add_gate_slice(index):\n",
" gate_slice_layer = network.add_slice(input=gates_output,\n",
" start=(0, index * hidden_size), # Start [batch_idx=0, col_idx]\n",
" shape=(batch_size, hidden_size), # Slice shape\n",
" stride=(1, 1)) # Stride\n",
" return gate_slice_layer.get_output(0)\n",
"\n",
" slice_i = add_gate_slice(0)\n",
" slice_i.name = \"Slice_I\"\n",
" slice_f = add_gate_slice(1)\n",
" slice_f.name = \"Slice_F\"\n",
" slice_c = add_gate_slice(2)\n",
" slice_c.name = \"Slice_C_candidate\" # Cell candidate\n",
" slice_o = add_gate_slice(3)\n",
" slice_o.name = \"Slice_O\"\n",
"\n",
" # Apply activations\n",
" act_i_layer = network.add_activation(slice_i, trt.ActivationType.SIGMOID)\n",
" act_i = act_i_layer.get_output(0)\n",
" act_i.name = \"Gate_I\"\n",
" act_f_layer = network.add_activation(slice_f, trt.ActivationType.SIGMOID)\n",
" act_f = act_f_layer.get_output(0)\n",
" act_f.name = \"Gate_F\"\n",
" act_c_layer = network.add_activation(slice_c, trt.ActivationType.TANH)\n",
" act_c = act_c_layer.get_output(0)\n",
" act_c.name = \"Gate_C_candidate\"\n",
" act_o_layer = network.add_activation(slice_o, trt.ActivationType.SIGMOID)\n",
" act_o = act_o_layer.get_output(0)\n",
" act_o.name = \"Gate_O\"\n",
"\n",
" # Cell state update: c_t = f_t * c_{t-1} + i_t * g_t\n",
" term1_c = network.add_elementwise(act_f, prev_c, trt.ElementWiseOperation.PROD)\n",
" term2_c = network.add_elementwise(act_i, act_c, trt.ElementWiseOperation.PROD)\n",
" next_c_layer = network.add_elementwise(term1_c.get_output(0), term2_c.get_output(0), trt.ElementWiseOperation.SUM)\n",
"\n",
" next_c = next_c_layer.get_output(0)\n",
" next_c.name = \"next_c\" # Shape [batch, hidden]\n",
"\n",
" # Hidden state update: h_t = o_t * tanh(c_t)\n",
" tanh_c_layer = network.add_activation(next_c, trt.ActivationType.TANH)\n",
" tanh_c = tanh_c_layer.get_output(0)\n",
" next_h_layer = network.add_elementwise(act_o, tanh_c, trt.ElementWiseOperation.PROD)\n",
"\n",
" next_h = next_h_layer.get_output(0)\n",
" next_h.name = \"next_h\" # Shape [batch, hidden]\n",
"\n",
" return next_h, next_c\n",
"\n",
"\n",
"def add_lstm_layer(network: trt.INetworkDefinition,\n",
" input_sequence: trt.ITensor, # Shape: [seq_len, batch_size, input_size]\n",
" hidden_size: int,\n",
" seq_len: int,\n",
" weight_W: np.ndarray, # [input_size, 4*hidden] or [hidden, 4*hidden]\n",
" weight_U: np.ndarray, # [hidden, 4*hidden]\n",
" bias: np.ndarray # [4*hidden]\n",
" ) -> trt.ITensor:\n",
" \"\"\"\n",
" Adds a LSTM to the network by adding one lstm_unit, and run multiple times with loops.\n",
" \"\"\"\n",
" # Infer batch_size and input_size from the input tensor shape\n",
" assert len(input_sequence.shape) == 3, f\"Input sequence tensor must have 3 dimensions [seq, batch, input]. Got shape {input_sequence.shape}\"\n",
" input_size = input_sequence.shape[2]\n",
"\n",
" # Shape: [batch_size, hidden_size]\n",
" initial_h = network.add_constant(np_initial_h.shape, np_initial_h).get_output(0)\n",
" initial_h.name = \"Initial_H\"\n",
" initial_c = network.add_constant(np_initial_c.shape, np_initial_c).get_output(0)\n",
" initial_c.name = \"Initial_C\"\n",
"\n",
" loop = network.add_loop()\n",
" loop.name = \"Time_Loop_Layer\"\n",
"\n",
" # add_trip_limit determines when the loop should stop. For here we want the loop to run seq_len times.\n",
" trip_limit = network.add_constant((), np.array([seq_len], dtype=np.int32)).get_output(0)\n",
" loop.add_trip_limit(trip_limit, trt.TripLimit.COUNT)\n",
"\n",
" # Recurrences for hidden and cell states\n",
" h_recurrence = loop.add_recurrence(initial_h)\n",
" c_recurrence = loop.add_recurrence(initial_c)\n",
" prev_h_tensor = h_recurrence.get_output(0)\n",
" prev_h_tensor.name = \"Prev_H\"\n",
" prev_c_tensor = c_recurrence.get_output(0)\n",
" prev_c_tensor.name = \"Prev_C\"\n",
"\n",
" # add_iterator iterates through slices of the input sequence along the specified axis, providing one slice per iteration.\n",
" x_t_iterator = loop.add_iterator(input_sequence, axis=0)\n",
" x_t = x_t_iterator.get_output(0)\n",
" x_t.name = \"x_t\"\n",
"\n",
"\n",
" # Call the LSTM unit function\n",
" next_h, next_c = add_lstm_unit(network=network,\n",
" input_x=x_t,\n",
" prev_h=prev_h_tensor,\n",
" prev_c=prev_c_tensor,\n",
" W=weight_W,\n",
" U=weight_U,\n",
" bias=bias,\n",
" hidden_size=hidden_size,\n",
" input_size=input_size)\n",
"\n",
" # Feed the computed states back into the recurrence inputs\n",
" h_recurrence.set_input(1, next_h)\n",
" c_recurrence.set_input(1, next_c)\n",
"\n",
" # add_loop_output() collects the values in the loop and outputs them. For this example, we concatenate the values along the first axis.\n",
" loop_output_h = loop.add_loop_output(next_h, trt.LoopOutput.CONCATENATE, axis=0)\n",
"\n",
" # when using CONCATENATE, the second input must be the trip limit.\n",
" loop_output_h.set_input(1, trip_limit)\n",
" loop_output_h.get_output(0).name = \"Hidden_Sequence\"\n",
"\n",
" # --- End of time step loop definition ---\n",
"\n",
" layer_output_sequence = loop_output_h.get_output(0)\n",
"\n",
" # The final output sequence is the sequence from the last layer\n",
" if layer_output_sequence is None:\n",
" raise RuntimeError(\"LSTM layer output was not generated (num_layers may be 0)\")\n",
" layer_output_sequence.name = \"Final_LSTM_Output_Sequence\"\n",
" return layer_output_sequence"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Step 2: Build the Network\n",
"\n",
"Now that we have the LSTM layer implementation (`add_lstm_layer`), let's proceed to build the TensorRT `INetworkDefinition`.\n",
"This involves defining the network structure by:\n",
"1. Adding the input tensor using `network.add_input`.\n",
"2. Adding the LSTM layer using our custom `add_lstm_layer` function.\n",
"3. Marking the LSTM layer's output tensor as the network's final output."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"builder = trt.Builder(TRT_LOGGER)\n",
"network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.STRONGLY_TYPED))\n",
"\n",
"# === Network Definition ===\n",
"# Shape: [seq_len, batch_size, input_size] -> e.g., [5, 1, 1]\n",
"input_tensor = network.add_input(name='input', dtype=trt.float32, shape=(seq_len, batch_size, input_size))\n",
"\n",
"# --- Add SINGLE LSTM Layer ---\n",
"lstm_output = add_lstm_layer(network=network,\n",
" input_sequence=input_tensor,\n",
" hidden_size=hidden_size,\n",
" seq_len=seq_len,\n",
" weight_W=np_weight_W, \n",
" weight_U=np_weight_U, \n",
" bias=np_bias) \n",
"# lstm_output shape: [seq_len, batch_size, hidden_size] -> e.g., [5, 1, 2]\n",
"\n",
"# --- Mark Output ---\n",
"lstm_output.name = 'hidden_state_sequence'\n",
"network.mark_output(lstm_output)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Step 3: Build the Engine\n",
"\n",
"Now that we have defined the network (`INetworkDefinition`), the next step is to build the optimized TensorRT engine. This process involves using the `trt.Builder` along with an `trt.BuilderConfig` object to specify how the engine should be built.\n",
"\n",
"The `IBuilderConfig` allows you to control various aspects of the build process, such as:\n",
"* Setting memory constraints (e.g., workspace size using `set_memory_pool_limit`).\n",
"* Setting builder flags to control optimization strategies and compatibility.\n",
"\n",
"Once the network and configuration are ready, the `builder.build_serialized_network(network, config)` method is called to produce the serialized engine, which can then be saved to a file or used directly.\n"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## (Optional) Defining a Progress Monitor\n",
"Building a TensorRT engine can sometimes take a while, especially for complex models. Don't worry if the build seems long! TensorRT offers a helpful tool called `IProgressMonitor`. This interface lets you track the build process step-by-step, making it easier to monitor progress and even debug if needed. \n",
"\n",
"### Implementing `IProgressMonitor`\n",
"\n",
"To use the progress monitor, inherit from `trt.IProgressMonitor` and override its key methods:\n",
"\n",
"* `phase_start(self, phase_name, parent_phase, num_steps)`: TensorRT calls this method when it begins a significant phase of the build process (e.g., \"Parsing ONNX Model\", \"Building Engine\"). \n",
" * `phase_name`: Name of the phase starting.\n",
" * `parent_phase`: Name of the parent phase, if this is a sub-phase (can be `None`).\n",
" * `num_steps`: The total number of steps expected for this phase.\n",
"* `step_complete(self, phase_name, step)`: Called after each incremental step within a phase is completed.\n",
" * `phase_name`: Name of the current phase.\n",
" * `step`: The index of the step that just finished (0-based).\n",
" * *Your implementation* usually updates the corresponding progress indicator.\n",
" * **Crucially, this method must return `True` to allow the build to continue.** Returning `False` or `None` will signal TensorRT to cancel the build.\n",
"* `phase_finish(self, phase_name)`: Called when a phase (and all its steps) is completed.\n",
" * `phase_name`: Name of the phase that finished.\n",
" * *Your implementation* typically finalizes and removes the progress indicator for this phase.\n",
"\n",
"After that, hook it with `IBuilderConfig` by setting `config.progress_monitor = MyProgressMonitor()`"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"class SimpleProgressMonitor(trt.IProgressMonitor):\n",
" def __init__(self):\n",
" trt.IProgressMonitor.__init__(self)\n",
" self._active_phases = 0\n",
"\n",
" def phase_start(self, phase_name, parent_phase, num_steps):\n",
" print(f\"[ProgressMonitor] Phase Start: {phase_name} ({num_steps} steps)\")\n",
" self._active_phases += 1\n",
"\n",
" def phase_finish(self, phase_name):\n",
" print(f\"[ProgressMonitor] Phase Finish: {phase_name}\")\n",
" self._active_phases -= 1\n",
"\n",
" def step_complete(self, phase_name, step):\n",
" print(f\"[ProgressMonitor] Step Complete: {phase_name}, Step {step}\")\n",
" return True\n",
"\n",
" @property\n",
" def active_phases(self):\n",
" return self._active_phases"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## (Optional) Version Compatible Engine\n",
"TensorRT engines are typically optimized for the specific GPU and TensorRT version they are built on. This maximizes performance but can cause incompatibility if the deployment environment differs.\n",
"\n",
"The `trt.BuilderFlag.VERSION_COMPATIBLE` flag addresses this by creating a more portable engine. This engine is less sensitive to minor variations in TensorRT versions or GPU models (within a compatible family), potentially at the cost of some performance compared to a non-compatible engine optimized for the exact target. It also reduces the need to rebuild the engine for every minor TensorRT update. Version compatibility is supported from TensorRT 8.6 onwards; the plan must be built with a version at least 8.6 or higher, and the runtime must also be 8.6 or higher.\n",
"\n",
"### Use Cases\n",
"* Deploying across diverse hardware fleets with compatible GPUs/TRT versions.\n",
"* Distributing applications where end-user system configurations vary.\n",
"* Simplifying maintenance by avoiding frequent rebuilds for minor updates.\n",
"\n",
"### How it Works\n",
"Enabling `trt.BuilderFlag.VERSION_COMPATIBLE` instructs TensorRT to use more generic optimizations. By default, this flag also causes a copy of a \"lean runtime\" (a version-specific, stripped-down runtime component) to be packaged within the engine plan file. When you deserialize this engine plan on a compatible system, TensorRT recognizes the embedded lean runtime, loads it, and uses this runtime to deserialize and execute the rest of the plan. \n",
"\n",
"Because this process involves loading and executing code (the lean runtime) directly from the engine plan file, you must explicitly indicate that you trust the origin and integrity of the plan. This is done by setting `runtime.engine_host_code_allowed = True` on your `trt.Runtime` instance before attempting to deserialize the engine.\n",
"\n",
"> **Considerations for Multiple Version-Compatible Engines:**\n",
"If deploying many version-compatible engines, the embedded lean runtime in each plan can lead to large overall application sizes. An alternative is to exclude the runtime from the engine plan (using `trt.BuilderFlag.EXCLUDE_LEAN_RUNTIME`) and load it manually. This approach can significantly reduce the total deployment footprint. For detailed instructions, refer to the NVIDIA TensorRT documentation on [Manually Loading the Runtime](https://docs.nvidia.com/deeplearning/tensorrt/latest/inference-library/advanced.html#manually-loading-the-runtime)."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"ENGINE_FILE_PATH = './lstm_network.plan'\n",
"\n",
"config = builder.create_builder_config()\n",
"config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 1 << 28) # 256MB\n",
"config.progress_monitor = SimpleProgressMonitor()\n",
"config.set_flag(trt.BuilderFlag.VERSION_COMPATIBLE)\n",
"\n",
"print(\"Building engine...\")\n",
"serialized_engine = builder.build_serialized_network(network, config)\n",
"\n",
"print(\"Engine build completed.\")\n",
"with open(ENGINE_FILE_PATH, 'wb') as f:\n",
" f.write(serialized_engine)\n",
"print(f\"Engine saved to {ENGINE_FILE_PATH}\")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Inference\n",
"\n",
"Once the TensorRT engine is built, the next step is typically to run inference to verify its functionality and performance. The standard process involves creating a runtime and execution context, managing GPU memory for inputs and outputs, transferring data between host and device, and executing the engine etc. While this process provides fine-grained control, it involves boilerplate code. This standard procedure was demonstrated in detail in Sample 1.\n",
"\n",
"In this sample, we'll simplify the inference process by using **[Polygraphy](https://github.com/NVIDIA/TensorRT/tree/main/tools/Polygraphy)**, a versatile toolkit included with TensorRT that automates many underlying details, such as:\n",
"* Context creation\n",
"* Buffer management\n",
"* Data transfers\n",
"\n",
"> **Important Note:** While Polygraphy is excellent for debugging and testing due to its ease of use, it may introduce overhead.\n",
"> For optimal performance in deployment scenarios, consider handcrafting the inference code as demonstrated in the `1_run_onnx_with_tensorrt` sample.\n",
"\n",
"For more examples, please refer to [polygraphy examples](https://github.com/NVIDIA/TensorRT/tree/main/tools/Polygraphy/examples)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"from polygraphy.backend.common import BytesFromPath\n",
"from polygraphy.backend.trt import EngineFromBytes, TrtRunner\n",
"\n",
"def run_inference_with_polygraphy(h_input: np.ndarray) -> np.ndarray:\n",
" input_name = 'input'\n",
" output_name = 'hidden_state_sequence'\n",
"\n",
" # Prepare the feed dictionary required by Polygraphy\n",
" # Ensure input is contiguous C-style array, which Polygraphy prefers.\n",
" h_input_contiguous = np.ascontiguousarray(h_input)\n",
" feed_dict = {input_name: h_input_contiguous}\n",
"\n",
" print(f\"Loading engine from: {ENGINE_FILE_PATH}\")\n",
" outputs = None\n",
" load_engine = EngineFromBytes(BytesFromPath(ENGINE_FILE_PATH))\n",
" with TrtRunner(load_engine) as runner:\n",
" outputs = runner.infer(feed_dict=feed_dict)\n",
" # Polygraphy automatically synchronizes, so no explicit stream sync needed here\n",
"\n",
" output_sequence = outputs[output_name]\n",
" print(f\"Output '{output_name}' shape: {output_sequence.shape}, dtype: {output_sequence.dtype}\")\n",
" return output_sequence\n",
"\n",
"output_sequence = run_inference_with_polygraphy(np_inputs) \n",
"\n",
"if output_sequence is not None:\n",
" print(f\"\\nInput Sequence (shape {np_inputs.shape}):\\n{np_inputs}\")\n",
" print(f\"\\nOutput Hidden State Sequence (shape {output_sequence.shape}):\\n{output_sequence}\")\n",
"else:\n",
" print(\"Inference failed.\")\n",
" "
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Verifying the Output (Comparison with Equivalent Operations in NumPy)\n",
"\n",
"To ensure our TensorRT LSTM implementation is correct, we'll compare its output with a reference implementation in NumPy. This is a common practice to validate custom layer logic.\n",
"\n",
"The NumPy version will mimic the same LSTM cell computations and unroll the loop over the time sequence."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"def sigmoid_np(x):\n",
" x_clipped = np.clip(x, -500, 500) # avoid overflow\n",
" return 1.0 / (1.0 + np.exp(-x_clipped))\n",
"\n",
"\n",
"def tanh_np(x):\n",
" x_clipped = np.clip(x, -100, 100) # avoid overflow\n",
" return np.tanh(x_clipped)\n",
"\n",
"\n",
"def lstm_step_numpy(x_t, prev_h, prev_c, W, U, bias):\n",
" # W: shape [input_size, 4*hidden_size]\n",
" # U: shape [hidden_size, 4*hidden_size]\n",
" # bias: shape [4*hidden_size]\n",
" # x_t: shape [batch_size, input_size]\n",
" # prev_h, prev_c: shape [batch_size, hidden_size]\n",
"\n",
" hidden_size_ = prev_h.shape[1]\n",
"\n",
" Wx = x_t @ W # Shape [batch_size, 4*hidden_size]\n",
" Uh = prev_h @ U # Shape [batch_size, 4*hidden_size]\n",
" gates = Wx + Uh + bias\n",
"\n",
" # Split gates\n",
" i = gates[:, 0 * hidden_size_ : 1 * hidden_size_]\n",
" f = gates[:, 1 * hidden_size_ : 2 * hidden_size_]\n",
" c = gates[:, 2 * hidden_size_ : 3 * hidden_size_] # Cell candidate\n",
" o = gates[:, 3 * hidden_size_ : 4 * hidden_size_]\n",
"\n",
" i_act = sigmoid_np(i)\n",
" f_act = sigmoid_np(f)\n",
" c_act = tanh_np(c)\n",
" o_act = sigmoid_np(o)\n",
"\n",
" next_c = f_act * prev_c + i_act * c_act\n",
" next_h = o_act * tanh_np(next_c)\n",
"\n",
" return next_h, next_c\n",
"\n",
"\n",
"def lstm_layer_numpy(input_sequence_np, np_W, np_U, np_bias):\n",
" seq_len_ = input_sequence_np.shape[0]\n",
" final_output_sequence_np = None\n",
" h = np_initial_h.copy()\n",
" c = np_initial_c.copy()\n",
"\n",
" layer_output_sequence_list = []\n",
"\n",
" for t in range(seq_len_):\n",
" # Slice the input sequence for this time step\n",
" x_t = input_sequence_np[t, :, :]\n",
"\n",
" h, c = lstm_step_numpy(x_t, h, c, np_W, np_U, np_bias)\n",
" layer_output_sequence_list.append(h)\n",
" layer_output_sequence_np = np.stack(layer_output_sequence_list, axis=0)\n",
" final_output_sequence_np = layer_output_sequence_np\n",
"\n",
" return final_output_sequence_np\n",
"\n",
"\n",
"numpy_output_sequence = lstm_layer_numpy(np_inputs, np_weight_W, np_weight_U, np_bias)\n",
"print(\"\\n--- NumPy LSTM Calculation Results ---\")\n",
"print(f\"Input Sequence (all ones, shape {np_inputs.shape}):\\n{np_inputs}\")\n",
"print(f\"\\nNumPy Output Hidden State Sequence (shape {numpy_output_sequence.shape}):\\n{numpy_output_sequence}\")\n",
"print(\"\\n--- Comparison ---\")\n",
"diff = np.abs(output_sequence - numpy_output_sequence)\n",
"max_diff = np.max(diff) if diff.size > 0 else 0.0\n",
"print(f\"Max absolute difference: {max_diff}\")\n",
"assert np.allclose(\n",
" output_sequence, numpy_output_sequence, atol=1e-5\n",
"), f\"Output sequence mismatch between TensorRT and NumPy, max diff: {max_diff}\"\n",
"print(\"Notebook executed successfully\")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Conclusion and Next Steps\n",
"\n",
"Congratulations! You have successfully:\n",
"- Defined an LSTM cell and layer using TensorRT's Layer APIs.\n",
"- Implemented a recurrent loop with `add_loop`.\n",
"- Monitored the engine build process using `IProgressMonitor`.\n",
"- Built a version-compatible TensorRT engine.\n",
"- Performed inference using the built engine via Polygraphy.\n",
"- Verified the results against a NumPy implementation.\n",
"\n",
"This sample demonstrates the fundamental building blocks for creating custom network architectures in TensorRT. From here, you can explore:\n",
"- More complex network structures.\n",
"- Different types of layers available in the TensorRT API.\n",
"- Advanced loop constructs and conditional logic.\n",
"- Further optimization techniques if performance is critical (though for this sample, we focused on API usage).\n",
"\n",
"By mastering the Layer API, you gain the power to optimize virtually any deep learning model for inference on NVIDIA GPUs."
]
}
],
"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
}
samples/python/refactored/README.md
+30
View File
@@ -0,0 +1,30 @@
# TensorRT Refactored Samples
This directory contains refactored and improved versions of TensorRT samples, demonstrating best practices and modern implementations.
## Available Samples
| Sample Name | Description | Format |
|-------------|-------------|---------|
| [1_run_onnx_with_tensorrt](./1_run_onnx_with_tensorrt) | Demonstrates ONNX model conversion to TensorRT and inference comparison | `ipynb` |
| [2_construct_network_with_layer_apis](./2_construct_network_with_layer_apis) | Constructing a Network with TensorRT Layer APIs | `ipynb` |
## Launch Instructions
1. Navigate to the desired sample directory and start the Jupyter server:
```bash
pip install notebook
cd 1_run_onnx_with_tensorrt # or any other sample
jupyter notebook
```
2. Then, open the `main.ipynb` file in the Jupyter Notebook interface that opens in your web browser.
# Changelog
October 2025
Migrate to strongly typed APIs.
August 2025
Removed support for Python versions < 3.10.