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) Has been cancelled
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
quickstart/IntroNotebooks/0. Running This Guide.ipynb
+144
View File
@@ -0,0 +1,144 @@
{
"cells": [
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# Running This Guide:"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"This guide is presented as a series of Jupyter notebooks covering both Tensorflow and PyTorch using a Python runtime.\n",
"\n",
"If you would like to run this code yourself, you can do so using the following steps:"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"__Setup__:\n",
"\n",
"The basic prerequisites you will need to use TensorRT are:\n",
"\n",
"- a [supported NVIDIA GPU](https://docs.nvidia.com/deeplearning/tensorrt/support-matrix/index.html#hardware-precision-matrix) (preferred compute capability 7.0 or higher - which includes INT8 precision support) \n",
"- latest [NVIDIA GPU drivers](https://docs.nvidia.com/datacenter/tesla/tesla-installation-notes/index.html)\n",
"- a [supported CUDA and cuDNN](https://docs.nvidia.com/deeplearning/tensorrt/support-matrix/index.html#platform-matrix) installation\n",
"\n",
"You can make sure your GPU environment is properly configured and check which GPU and CUDA version you are using with nvidia-smi:"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"!nvidia-smi "
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"For some of the examples you will also need cuda-python, skimage, matplotlib, and onnx:"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"%pip install cuda-python onnx scikit-image matplotlib"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"__PyTorch__:\n"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"We will be using PyTorch to walk through the basic steps of deploying a TensorRT model by [exporting the model in ONNX format](https://pytorch.org/docs/stable/onnx.html).\n",
"\n",
"You can find PyTorch installation instructions [here](https://pytorch.org/get-started/locally/), or use one of NVIDIA's NGC containers [here](https://ngc.nvidia.com/catalog/containers/nvidia:pytorch). \n",
"\n",
"You will also need torchvision:"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"%pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu124"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"__Using Colab:__\n",
"\n",
"You can also test these notebooks out using [Google Colab](https://colab.research.google.com/), which includes PyTorch, as well as supported NVIDIA drivers. __Make sure to select a GPU hardware accelerator__ in the runtime options. Just note that TensorRT performance is best on newer gpus, Colab often has trouble with reduced precision inference, and you will have to use an older version of TensorRT."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"__TensorRT Support:__\n",
"\n",
"TensorRT support for NVIDIA GPUs is determined by their __compute capability__. You can check the compute cabapility of your card on the [NVIDIA website](https://developer.nvidia.com/cuda-gpus).\n",
"\n",
"TensorRT supports different feautures depending on your compute capability. Higher compute capabilities allow additional TensorRT optimizations, like reduced precision inference. You can check which TensorRT features are supported by your compute capability in the [TensorRT documentation](https://docs.nvidia.com/deeplearning/tensorrt/support-matrix/index.html#hardware-precision-matrix).\n",
"\n",
"\n"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"__Next Steps:__\n",
"\n",
"Now, start by opening [1. Introduction.ipynb](./1.%20Introduction.ipynb) and proceed through the notebook!\n",
"\n",
"The notebooks included with this guide are:\n",
"- [1. Introduction.ipynb](./1.%20Introduction.ipynb)\n",
"- [2. Using PyTorch through ONNX.ipynb](./2.%20Using%20PyTorch%20through%20ONNX.ipynb)\n",
"- [3. Understanding TensorRT Runtimes.ipynb](./3.%20Understanding%20TensorRT%20Runtimes.ipynb)"
]
}
],
"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.10.12"
}
},
"nbformat": 4,
"nbformat_minor": 4
}
quickstart/IntroNotebooks/1. Introduction.ipynb
+392
View File
@@ -0,0 +1,392 @@
{
"cells": [
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# Getting Started with TensorRT"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"TensorRT is an SDK for optimizing trained deep learning models to enable high-performance inference. TensorRT contains a deep learning inference __optimizer__ for trained deep learning models and an optimized __runtime__ for execution. After you have trained your deep learning model in a framework of your choice, TensorRT enables you to run it with higher throughput and lower latency. \n",
"\n",
"The TensorRT ecosystem breaks broadly down into two parts:\n",
"<br><br><br>\n",
"![TensorRT Landscape](./images/tensorrt_landscape.png)\n",
"<br><br><br>\n",
"Essentially,\n",
"\n",
"1. The various paths users can follow to convert their models to optimized TensorRT engines\n",
"2. The various runtimes users can target with TensorRT when deploying their optimized TensorRT engines\n",
"\n",
"If you have a model in PyTorch and want to run inference as efficiently as possible - with low latency, high throughput, and less memory consumption - this guide will help you achieve just that!!"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## How Do I Use TensorRT:"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"TensorRT is a large and flexible project. It can handle a variety of workflows, and which workflow is best for you will depend on your specific use-case and problem setting. Abstractly, the process for deploying a model from a deep learning framework to TensorRT looks like this:\n",
"\n",
"![TensorRT Workflow](./images/tensorrt_workflow.png)\n",
"\n",
"To help you get there, this guide will help you answer five key questions:\n",
"\n",
"1. __What format should I save my model in?__\n",
"2. __What batch size(s) am I running inference at?__\n",
"3. __What precision am I running inference at?__\n",
"4. __What TensorRT path am I using to convert my model?__\n",
"5. __What runtime am I targeting?__\n",
"\n",
"This guide will walk you broadly through all of these decision points while giving you an overview of your options at each step.\n",
"\n",
"We could talk about these points in isolation, but they are best understood in the context of an actual end-to-end workflow. Let's get started on a simple one here, using a TensorRT API wrapper written for this guide. Once you understand the basic workflow, you can dive into the more in depth notebooks on the Torch-TRT and ONNX converters!"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Simple TensorRT Demonstration through ONNX:"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"There are several ways of approaching TensorRT conversion and deployment. Here, we will take a pretrained ResNet50 model, convert it to an optimized TensorRT engine, and run it in the TensorRT runtime.\n",
"\n",
"For this simple demonstration we will focus the ONNX path - one of the two main automatic approaches for TensorRT conversion. We will then run the model in the TensorRT Python API using a simplified wrapper written for this guide. Essentially, we will follow this path to convert and deploy our model:\n",
"\n",
"![ONNX Conversion](./images/onnx_onnx.png)\n",
"\n",
"We will follow the five questions above. For a more in depth discussion, the section following this demonstration will cover options available at these steps in more detail."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"__IMPORTANT NOTE:__ Please __shutdown all other notebooks and PyTorch processes__ before running these steps. TensorRT and PyTorch can not be loaded into your Python processes at the same time."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"#### 1. What format should I save my model in?"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"The two main automatic conversion paths for TensorRT require different model formats to successfully convert a model. Torch-TRT uses PyTorch saved models and the ONNX path requires models be saved in ONNX. Here, we will use ONNX.\n",
"\n",
"We are going to use ResNet50 - a basic backbone vision model that can be used for a variety of purposes. For the sake of demonstration, here we will perform classification using a __pretrained ResNet50 ONNX__ model included with the [ONNX model zoo](https://github.com/onnx/models).\n",
"\n",
"We can download a pretrained ResNet50 from the ONNX model zoo and untar it by doing the following:"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"!wget https://download.onnxruntime.ai/onnx/models/resnet50.tar.gz -O resnet50.tar.gz\n",
"!tar xzf resnet50.tar.gz"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"See how to export ONNX models that will work with this same trtexec command in the [PyTorch through ONNX notebook](./2.%20Using%20PyTorch%20through%20ONNX.ipynb)."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"#### 2. What precision will I use?"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Inference typically requires less numeric precision than training. With some care, lower precision can give you faster computation and lower memory consumption without sacrificing any meaningful accuracy. TensorRT supports TF32, FP32, FP16, and INT8 precisions.\n",
"\n",
"FP32 is the default training precision of most frameworks, so we will start by using FP32 for inference here. Let's create a \"dummy\" batch to work with in order to test our model. TensorRT will use the precision of the input batch throughout the rest of the network by default."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"import numpy as np\n",
"PRECISION = np.float32\n",
"\n",
"# The input tensor shape of the ONNX model.\n",
"input_shape = (1, 3, 224, 224)\n",
"\n",
"dummy_input_batch = np.zeros(input_shape, dtype=PRECISION)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"#### 3. What TensorRT path am I using to convert my model?"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"The ONNX conversion path is one of the most universal and performant paths for automatic TensorRT conversion. It works for PyTorch and many other frameworks. There are several tools to help users convert models from ONNX to a TensorRT engine. \n",
"\n",
"One common approach is to use trtexec - a command line tool included with TensorRT that can, among other things, convert ONNX models to TensorRT engines and profile them."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"!trtexec --onnx=resnet50/model.onnx --saveEngine=resnet_engine_intro.engine "
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"__Notes on the flags above:__\n",
" \n",
"Tell trtexec where to find our ONNX model:\n",
"\n",
" --onnx=resnet50/model.onnx \n",
"\n",
"Tell trtexec where to save our optimized TensorRT engine:\n",
"\n",
" --saveEngine=resnet_engine_intro.engine\n"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"#### 4. What runtime will I use?\n",
"\n",
"After we have our TensorRT engine created successfully, we need to decide how to run it with TensorRT.\n",
"\n",
"There are two types of TensorRT runtimes: a standalone runtime which has C++ and Python bindings, and a native integration into PyTorch. In this section, we will use a simplified wrapper (ONNXClassifierWrapper) which calls the standalone runtime. "
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"!python3 -m pip install \"numpy<2.0\""
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# If you get an error in this cell, restart your notebook (possibly your whole machine) and do not run anything that imports/uses PyTorch\n",
"\n",
"from onnx_helper import ONNXClassifierWrapper\n",
"trt_model = ONNXClassifierWrapper(\"resnet_engine_intro.engine\", target_dtype = PRECISION)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"__Note__: If this conversion fails, please restart your Jupyter notebook kernel (in menu bar Kernel->Restart Kernel) and run steps 3 to 5 again. If you get an error like 'TypeError: pybind11::init(): factory function returned nullptr' there is likely some dangling process on the GPU - restart your machine and try again.\n",
"\n",
"We will feed our batch of randomized dummy data into our ONNXClassifierWrapper to run inference on that batch:"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# Warm up:\n",
"trt_model.predict(dummy_input_batch)[:10] # softmax probability predictions for the first 10 classes of the first sample"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"We can get a rough sense of performance using %%timeit:"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"%%timeit\n",
"trt_model.predict(dummy_input_batch)[:10]"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Applying TensorRT to Your Model:"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"This is a simple example applied to a single model, but how should you go about answering these questions for your workload?\n",
"\n",
"First and foremost, it is a good idea to get an understanding of what your options are, and where you can learn more about them! \n",
"\n",
"### __Compatible Models:__ MLP/CNN/RNN/Transformer/Embedding/Etc\n",
"\n",
"TensorRT is compatible with models consisting of [these layers](https://docs.nvidia.com/deeplearning/tensorrt/support-matrix/index.html#layers-matrix). Using only supported layers ensures optimal performance without having to write any custom plugin code.\n",
"\n",
"In terms of framework, TensorRT is integrated directly with PyTorch - and most other major deep learning frameworks are supported by first converting to ONNX format.\n",
"\n",
"### __Conversion Methods:__ ONNX/TensorRT API\n",
"\n",
"The __ONNX__ path is the most performant and framework-agnostic automatic way of converting models. It's main disadvantage is that it must convert networks completely - if a network has an unsupported layer ONNX can't convert it unless you write a custom plugin.\n",
"\n",
"You can see an example of how to use TensorRT with ONNX:\n",
"- [Here](./2.%20Using%20PyTorch%20through%20ONNX.ipynb) in this guide for PyTorch\n",
"\n",
"There is the __TensorRT API__. The TensorRT ONNX path automatically convert models to TensorRT engines for you. Sometimes, however, we want to convert something complex, or have the maximum amount of control in how our TensorRT engine is created. This let's us do things like creating custom plugins for layers that TensorRT doesn't support. \n",
"\n",
"When using this approach, we create TensorRT engine manually operation-by-operation using the TensorRT API's available in Python and C++. This process involves building a network identical in structure to your target network using the TensorRT API, and then loading in the weights directly in proper format. You can find more details on this [in the TensorRT documentation](https://docs.nvidia.com/deeplearning/tensorrt/developer-guide/index.html#c_topics).\n",
"\n",
"### __Precision:__ TF32/FP32/FP16/INT8/FP8\n",
"\n",
"TensorRT feature support - such as precision - for NVIDIA GPUs is determined by their __compute capability__. You can check the compute cabapility of your card on the [NVIDIA website](https://developer.nvidia.com/cuda-gpus).\n",
"\n",
"TensorRT supports different precisions depending on said compute capability. You can check what features are supported by your compute capability in the [TensorRT documentation](https://docs.nvidia.com/deeplearning/tensorrt/support-matrix/index.html#hardware-precision-matrix).\n",
"\n",
"__TF32__ is the default training precision on cards with compute cabapilities 8.0 and higher (e.g. NVIDIA A100 and later) - use when you want to replicate your original model performance as closely as possible on cards with compute capability of 8.0 or higher. \n",
"\n",
"TF32 is a precision designed to preserve the range of FP32 with the precision of FP16. In practice, this means that TF32 models train faster than FP32 models while still converging to the same accuracy. This feature is only available on newer GPUs.\n",
"\n",
"__FP32__ is the default training precision on cards with compute cabapilities of less than 8.0 (e.g. pre-NVIDIA A100) - use when you want to replicate your original model performance as closely as possible on cards with compute capability of less than 8.0\n",
"\n",
"__FP16__ is an inference focused reduced precision. It gives up some accuracy for faster models with lower latency and lower memory footprint. In practice, the accuracy loss is generally negligible in FP16 - so FP16 is a fairly safe bet in most cases for inference. Cards that are focused on deep learning training often have strong FP16 capabilities, making FP16 a great choice for GPUs that are expected to be used for both training and inference.\n",
"\n",
"__INT8__ is an inference focused reduced precision. It further reduces memory requirements and latency compared to FP16. INT8 has the potential to lose more accuracy than FP16 - but TensorRT provides tools to help you quantize your network's INT8 weights to avoid this as much as possible. INT8 requires the extra step of calibrating how TensorRT should quantize your weights to integers - requiring some sample data. With careful tuning and a good calibration dataset, accuracy loss from INT8 is often minimal. This makes INT8 a great precision for lower-power environments such as those using T4 GPUs or AGX Jetson modules - both of which have strong INT8 capabilities.\n",
"\n",
"__FP8__ 8-bit floating point type with 1-bit for sign, 4-bits for exponent, 3-bits for mantissa is an inference focused reduced precision. Similar to INT8, it further reduces memory requirements and latency compared to FP16 with potential suffer from precision loss, especially in deep models with large ranges of values. Extra steps of calibration are needed to minimize accuracy loss. The FP8 precision is supported on NVIDIA GPUs with compute cabapilities of 9.0 and higher (e.g. NVIDIA H100 and later). \n",
"\n",
"### __Runtime:__ Torch-TRT/Python API/C++ API/TRITON\n",
"\n",
"For a more in depth discussion of these options and how they compare see [this notebook on TensorRT Runtimes!](./3.%20Understanding%20TensorRT%20Runtimes.ipynb)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## What do I do if I run into issues with conversion?"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Here are several steps you can try if your model is not converting to TensorRT properly:\n",
"\n",
"1. Check the logs - if you are using a tool such as trtexec to convert your model, it will tell you which layer is problematic\n",
"2. Write a custom plugin - you can find more information on it [here](https://docs.nvidia.com/deeplearning/tensorrt/developer-guide/index.html#extending).\n",
"3. Use alternative implementations of the layers or operations in question in your network definition - for example, it can be easier to use the padding argument in your convolutional layers instead of adding an explicit padding layer to the network. \n",
"4. ONNX to TensorRT conversion can be hard to debug, but tools like graph surgeon https://docs.nvidia.com/deeplearning/tensorrt/api/python_api/graphsurgeon/graphsurgeon.html can help you fix specific nodes in your graph as well as pull it apart for analysis or patch specific nodes in your graph\n",
"5. Ask on the [NVIDIA developer forums](https://forums.developer.nvidia.com/c/ai-data-science/deep-learning/tensorrt) - we have many active TensorRT experts at NVIDIA who who browse the forums and can help\n",
"6. Post an issue on the [TensorRT OSS Github](https://github.com/NVIDIA/TensorRT)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Next Steps:\n",
"\n",
"You have now taken a model saved in ONNX format, converted it to an optimized TensorRT engine, and deployed it using the Python runtime. This is a great first step towards getting better performance out of your deep learning models at inference time!\n",
"\n",
"Now, you can check out the remaining notebooks in this guide. See:\n",
"\n",
"- [2. Using PyTorch through ONNX.ipynb](./2.%20Using%20PyTorch%20through%20ONNX.ipynb)\n",
"- [3. Understanding TensorRT Runtimes.ipynb](./3.%20Understanding%20TensorRT%20Runtimes.ipynb)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"<h4> Profiling </h4>\n",
"\n",
"This is a great next step for further optimizing and debugging models you are working on productionizing\n",
"\n",
"You can find it here: https://docs.nvidia.com/deeplearning/tensorrt/developer-guide/index.html#performance\n",
"\n",
"<h4> TRT Dev Docs </h4>\n",
"\n",
"Main documentation page for the ONNX, layer builder, C++, and legacy APIs\n",
"\n",
"You can find it here: https://docs.nvidia.com/deeplearning/tensorrt/developer-guide/index.html\n",
"\n",
"<h4> TRT OSS GitHub </h4>\n",
"\n",
"Contains OSS TRT components, sample applications, and plugin examples\n",
"\n",
"You can find it here: https://github.com/NVIDIA/TensorRT"
]
}
],
"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.10.12"
}
},
"nbformat": 4,
"nbformat_minor": 4
}
quickstart/IntroNotebooks/2. Using PyTorch through ONNX.ipynb
+665
View File
@@ -0,0 +1,665 @@
{
"cells": [
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# Using PyTorch with TensorRT through ONNX:"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"TensorRT is a great way to take a trained PyTorch model and optimize it to run more efficiently during inference on an NVIDIA GPU.\n",
"\n",
"One approach to convert a PyTorch model to TensorRT is to export a PyTorch model to ONNX (an open format exchange for deep learning models) and then convert into a TensorRT engine. Essentially, we will follow this path to convert and deploy our model:\n",
"\n",
"![PyTorch+ONNX](./images/pytorch_onnx.png)\n",
"\n",
"Both PyTorch and TensorFlow models can be exported to ONNX, as well as many other frameworks. This allows models created using either framework to flow into common downstream pipelines.\n",
"\n",
"To get started, let's take a well-known computer vision model and follow five key steps to deploy it to the TensorRT Python runtime:\n",
"\n",
"1. __What format should I save my model in?__\n",
"2. __What batch size(s) am I running inference at?__\n",
"3. __What precision am I running inference at?__\n",
"4. __What TensorRT path am I using to convert my model?__\n",
"5. __What runtime am I targeting?__"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## 1. What format should I save my model in?"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"We are going to use ResNet50, a widely used CNN architecture first described in <a href=https://arxiv.org/abs/1512.03385>this paper</a>.\n",
"\n",
"Let's start by loading dependencies and downloading the model. We will also move our Resnet model onto the GPU and set it to evaluation mode."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"import torchvision.models as models\n",
"import torch\n",
"import torch.onnx\n",
"\n",
"# load the pretrained model\n",
"resnet50 = models.resnet50(pretrained=True, progress=False).eval()"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"When saving a model to ONNX, PyTorch requires a test batch in proper shape and format. We pick a batch size:"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"BATCH_SIZE=32\n",
"\n",
"dummy_input=torch.randn(BATCH_SIZE, 3, 224, 224)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Next, we will export the model using the dummy input batch:"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# export the model to ONNX\n",
"torch.onnx.export(resnet50, dummy_input, \"resnet50_pytorch.onnx\", verbose=False)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Note that we are picking a BATCH_SIZE of 32 in this example."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Now Test with a Real Image:"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Let's try a real image batch! For this example, we will simply repeat one open-source dog image from http://www.dog.ceo:"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"from skimage import io\n",
"from skimage.transform import resize\n",
"from matplotlib import pyplot as plt\n",
"import numpy as np\n",
"\n",
"url='https://images.dog.ceo/breeds/retriever-golden/n02099601_3004.jpg'\n",
"img = resize(io.imread(url), (224, 224))\n",
"img = np.expand_dims(np.array(img, dtype=np.float32), axis=0) # Expand image to have a batch dimension\n",
"input_batch = np.array(np.repeat(img, BATCH_SIZE, axis=0), dtype=np.float32) # Repeat across the batch dimension\n",
"\n",
"input_batch.shape"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"plt.imshow(input_batch[0].astype(np.float32))"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"resnet50_gpu = models.resnet50(pretrained=True, progress=False).to(\"cuda\").eval()"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"We need to move our batch onto GPU and properly format it to shape [32, 3, 224, 224]. "
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"input_batch_chw = torch.from_numpy(input_batch).transpose(1,3).transpose(2,3)\n",
"input_batch_gpu = input_batch_chw.to(\"cuda\")\n",
"\n",
"input_batch_gpu.shape"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"We can run a prediction on a batch using .forward():"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"with torch.no_grad():\n",
" predictions = np.array(resnet50_gpu(input_batch_gpu).cpu())\n",
"\n",
"predictions.shape"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Verify Baseline Model Performance/Accuracy:"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"For a baseline, lets time our prediction in FP32:"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"%%timeit\n",
"\n",
"with torch.no_grad():\n",
" preds = np.array(resnet50_gpu(input_batch_gpu).cpu())"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"We can also time FP16 precision performance:"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"resnet50_gpu_half = resnet50_gpu.half()\n",
"input_half = input_batch_gpu.half()\n",
"\n",
"with torch.no_grad():\n",
" preds = np.array(resnet50_gpu_half(input_half).cpu()) # Warm Up\n",
"\n",
"preds.shape"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"%%timeit\n",
"\n",
"with torch.no_grad():\n",
" preds = np.array(resnet50_gpu_half(input_half).cpu())"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Let's also make sure our results are accurate. We will look at the top 5 accuracy on a single image prediction. The image we are using is of a Golden Retriever, which is class 207 in the ImageNet dataset our model was trained on."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"indices = (-predictions[0]).argsort()[:5]\n",
"print(\"Class | Likelihood\")\n",
"list(zip(indices, predictions[0][indices]))"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"We have a model exported to ONNX and a baseline to compare against! Let's now take our ONNX model and convert it to a TensorRT inference engine."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## 2. What batch size(s) am I running inference at?\n",
"\n",
"We are going to run with a fixed batch size of 32 for this example. Note that above we set BATCH_SIZE to 32 when saving our model to ONNX. We need to create another dummy batch of the same size (this time it will need to be in our target precision) to test out our engine.\n",
"\n",
"First, as before, we will set our BATCH_SIZE to 32. Note that our trtexec command above includes the '--explicitBatch' flag to signal to TensorRT that we will be using a fixed batch size at runtime."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"BATCH_SIZE = 32"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Importantly, by default TensorRT will use the input precision you give the runtime as the default precision for the rest of the network. So before we create our new dummy batch, we also need to choose a precision as in the next section:"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## 3. What precision am I running inference at?\n",
"\n",
"Remember that lower precisions than FP32 tend to run faster. There are two common reduced precision modes - FP16 and INT8. Graphics cards that are designed to do inference well often have an affinity for one of these two types. This guide was developed on an NVIDIA V100, which favors FP16, so we will use that here by default. INT8 is a more complicated process that requires a calibration step.\n",
"\n",
"__NOTE__: Make sure you use the same precision (USE_FP16) here you saved your model in above!"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"import numpy as np\n",
"\n",
"USE_FP16 = True\n",
"target_dtype = np.float16 if USE_FP16 else np.float32"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
" To create a test batch, we will once again repeat one open-source dog image from http://www.dog.ceo:"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"from skimage import io\n",
"from skimage.transform import resize\n",
"from matplotlib import pyplot as plt\n",
"import numpy as np\n",
"\n",
"url='https://images.dog.ceo/breeds/retriever-golden/n02099601_3004.jpg'\n",
"img = resize(io.imread(url), (224, 224))\n",
"input_batch = np.array(np.repeat(np.expand_dims(np.array(img, dtype=np.float32), axis=0), BATCH_SIZE, axis=0), dtype=np.float32)\n",
"\n",
"input_batch.shape"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"plt.imshow(input_batch[0].astype(np.float32))"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Preprocess Images:\n",
"\n",
"PyTorch has a normalization that it applies by default in all of its pretrained vision models - we can preprocess our images to match this normalization by the following, making sure our final result is in FP16 precision:"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"import torch\n",
"from torchvision.transforms import Normalize\n",
"\n",
"def preprocess_image(img):\n",
" norm = Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])\n",
" result = norm(torch.from_numpy(img).transpose(0,2).transpose(1,2))\n",
" return np.array(result, dtype=np.float16)\n",
"\n",
"preprocessed_images = np.array([preprocess_image(image) for image in input_batch])"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## 4. What TensorRT path am I using to convert my model?"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"We can use trtexec, a command line tool for working with TensorRT, in order to convert an ONNX model originally from PyTorch to an engine file.\n",
"\n",
"Let's make sure we have TensorRT installed (this comes with trtexec):"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"import tensorrt"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"To convert the model we saved in the previous step, we need to point to the ONNX file, give trtexec a name to save the engine as, and last specify that we want to use a fixed batch size instead of a dynamic one."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# step out of Python for a moment to convert the ONNX model to a TRT engine using trtexec\n",
"if USE_FP16:\n",
" !trtexec --onnx=resnet50_pytorch.onnx --saveEngine=resnet_engine_pytorch.trt --inputIOFormats=fp16:chw --outputIOFormats=fp16:chw --fp16\n",
"else:\n",
" !trtexec --onnx=resnet50_pytorch.onnx --saveEngine=resnet_engine_pytorch.trt"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"This will save our model as 'resnet_engine.trt'."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## 5. What TensorRT runtime am I targeting?\n",
"\n",
"Now, we have a converted our model to a TensorRT engine. Great! That means we are ready to load it into the native Python TensorRT runtime. This runtime strikes a balance between the ease of use of the high level Python APIs used in frameworks and the fast, low level C++ runtimes available in TensorRT."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"%%time\n",
"\n",
"import tensorrt as trt\n",
"from cuda.bindings import runtime as cudart\n",
"import numpy as np\n",
"\n",
"err, = cudart.cudaSetDevice(0)\n",
"assert err == cudart.cudaError_t.cudaSuccess\n",
"\n",
"f = open(\"resnet_engine_pytorch.trt\", \"rb\")\n",
"runtime = trt.Runtime(trt.Logger(trt.Logger.WARNING))\n",
"\n",
"engine = runtime.deserialize_cuda_engine(f.read())\n",
"context = engine.create_execution_context()"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Now allocate input and output memory, give TRT pointers (bindings) to it:"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"import numpy as np\n",
"\n",
"# need to set input and output precisions to FP16 to fully enable it\n",
"output = np.empty([BATCH_SIZE, 1000], dtype = target_dtype)\n",
"\n",
"# allocate device memory\n",
"err, d_input = cudart.cudaMalloc(input_batch.nbytes)\n",
"assert err == cudart.cudaError_t.cudaSuccess\n",
"err, d_output = cudart.cudaMalloc(output.nbytes)\n",
"assert err == cudart.cudaError_t.cudaSuccess\n",
"\n",
"tensor_names = [engine.get_tensor_name(i) for i in range(engine.num_io_tensors)]\n",
"assert(len(tensor_names) == 2)\n",
"\n",
"context.set_tensor_address(tensor_names[0], d_input)\n",
"context.set_tensor_address(tensor_names[1], d_output)\n",
"\n",
"err, stream = cudart.cudaStreamCreate()\n",
"assert err == cudart.cudaError_t.cudaSuccess"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Next, set up the prediction function.\n",
"\n",
"This involves a copy from CPU RAM to GPU VRAM, executing the model, then copying the results back from GPU VRAM to CPU RAM:"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"def predict(batch): # result gets copied into output\n",
" # transfer input data to device\n",
" err, = cudart.cudaMemcpyAsync(d_input, batch.ctypes.data, batch.nbytes,\n",
" cudart.cudaMemcpyKind.cudaMemcpyHostToDevice, stream)\n",
" assert err == cudart.cudaError_t.cudaSuccess\n",
"\n",
" # execute model\n",
" context.execute_async_v3(stream)\n",
"\n",
" # transfer predictions back\n",
" err, = cudart.cudaMemcpyAsync(output.ctypes.data, d_output, output.nbytes,\n",
" cudart.cudaMemcpyKind.cudaMemcpyDeviceToHost, stream)\n",
" assert err == cudart.cudaError_t.cudaSuccess\n",
"\n",
" # synchronize stream\n",
" err, = cudart.cudaStreamSynchronize(stream)\n",
" assert err == cudart.cudaError_t.cudaSuccess\n",
"\n",
" return output"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Let's time the function!"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"print(\"Warming up...\")\n",
"\n",
"pred = predict(preprocessed_images)\n",
"\n",
"print(\"Done warming up!\")"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"%%timeit\n",
"\n",
"pred = predict(preprocessed_images)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Finally we should verify our TensorRT output is still accurate."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"indices = (-pred[0]).argsort()[:5]\n",
"print(\"Class | Probability (out of 1)\")\n",
"list(zip(indices, pred[0][indices]))"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"err, = cudart.cudaStreamDestroy(stream)\n",
"assert err == cudart.cudaError_t.cudaSuccess\n",
"err, = cudart.cudaFree(d_input)\n",
"assert err == cudart.cudaError_t.cudaSuccess\n",
"err, = cudart.cudaFree(d_output)\n",
"assert err == cudart.cudaError_t.cudaSuccess"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Look for ImageNet indices 150-275 above, where 207 is the ground truth correct class (Golden Retriever). Compare with the results of the original unoptimized model in the first section!"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Next Steps:"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"<h4> TRT Dev Docs </h4>\n",
"\n",
"Main documentation page for the ONNX, layer builder, C++, and legacy APIs\n",
"\n",
"You can find it here: https://docs.nvidia.com/deeplearning/tensorrt/developer-guide/index.html\n",
"\n",
"<h4> TRT OSS GitHub </h4>\n",
"\n",
"Contains OSS TRT components, sample applications, and plugin examples\n",
"\n",
"You can find it here: https://github.com/NVIDIA/TensorRT\n",
"\n",
"\n",
"#### TRT Supported Layers:\n",
"\n",
"https://docs.nvidia.com/deeplearning/tensorrt/support-matrix/index.html#layers-precision-matrix\n",
"\n",
"#### TRT ONNX Plugin Example:\n",
"\n",
"https://github.com/NVIDIA/TensorRT/tree/main/samples/sampleOnnxMnistCoordConvAC"
]
}
],
"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.10.12"
}
},
"nbformat": 4,
"nbformat_minor": 5
}
quickstart/IntroNotebooks/3. Understanding TensorRT Runtimes.ipynb
+107
View File
@@ -0,0 +1,107 @@
{
"cells": [
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# Runtimes: What are my options? How do I choose?\n",
"\n",
"Remember that TensorRT consists of two main components - __1. A series of parsers and integrations__ to convert your model to an optimized engine and __2. An series of TensorRT runtime APIs__ with several associated tools for deployment.\n",
"\n",
"In this notebook, we will focus on the latter - various runtime options for TensorRT engines.\n",
"\n",
"The runtimes have different use cases for running TRT engines. \n",
"\n",
"### Considerations when picking a runtime:\n",
"\n",
"Generally speaking, there are a few major considerations when picking a runtime:\n",
"- __Framework__ - Some options, like Torch-TRT, are only relevant to PyTorch\n",
"- __Time-to-solution__ - Torch-TRT is much more likely to work 'out-of-the-box' if a quick solution is required and ONNX fails\n",
"- __Serving needs__ - Torch-TRT can use TorchServe to serve models over Cloud as a simple solution. For other frameworks (or for more advanced features) TRITON is framework agnostic, allows for concurrent model execution or multiple copies within a GPU to reduce latency, and can accept engines created through both the ONNX and Torch-TRT paths\n",
"- __Performance__ - Different TensorRT runtimes offer varying levels of performance. For example, Torch-TRT is generally going to be slower than using ONNX or the C++ API directly.\n",
"\n",
"### Python API:\n",
"\n",
"__Use this when:__\n",
"- You can accept some performance overhead, and\n",
"- You are most familiar with Python, or\n",
"- You are performing initial debugging and testing with TRT\n",
"\n",
"__More info:__\n",
"\n",
" \n",
"The [TensorRT Python API](https://docs.nvidia.com/deeplearning/tensorrt/developer-guide/index.html#perform_inference_python) gives you fine-grained control over the execution of your engine using a Python interface. It makes memory allocation, kernel execution, and copies to and from the GPU explicit - which can make integration into high performance applications easier. It is also great for testing models in a Python environment - such as in a Jupyter notebook.\n",
" \n",
"The [ONNX notebook for PyTorch](./2.%20Using%20PyTorch%20through%20ONNX.ipynb) is a good example of using TensorRT to get great performance while staying in Python.\n",
"\n",
"### C++ API: \n",
"\n",
"__Use this when:__\n",
"- You want the least amount of overhead possible to maximize the performance of your models and achieve better latency\n",
"- You are not using Torch-TRT (though Torch-TRT graph conversions that only generate a single engine can still be exported to C++)\n",
"- You are most familiar with C++\n",
"- You want to optimize your inference pipeline as much as possible\n",
"\n",
"__More info:__\n",
"\n",
"The [TensorRT C++ API](https://docs.nvidia.com/deeplearning/tensorrt/developer-guide/index.html#perform_inference_c) gives you fine-grained control over the execution of your engine using a C++ interface. It makes memory allocation, kernel execution, and copies to and from the GPU explicit - which can make integration into high performance C++ applications easier. The C++ API is generally the most performant option for running TensorRT engines, with the least overhead.\n",
"\n",
"[This NVIDIA Developer blog](https://developer.nvidia.com/blog/speed-up-inference-tensorrt/) is a good example of taking an ONNX model and running it with dynamic batch size support using the C++ API.\n",
"\n",
"\n",
"### Torch-TRT Runtime:\n",
" \n",
"__Use this when:__\n",
" \n",
"- You are using Torch-TRT, and\n",
"- Your model converts to more than one TensorRT engine\n",
"\n",
"__More info:__\n",
"\n",
"\n",
"Torch-TRT is the standard runtime used with models that were converted in Torch-TRT. It works by taking groups of nodes at once in the PyTorch graph, and replacing them with a singular optimized engine that calls the TensorRT Python API behind the scenes. This optimized engine is in the form of a PyTorch operation - which means that your graph is still in PyTorch and will essentially function like any other PyTorch model.\n",
"\n",
"If your graph entirely converts to a single Torch-TRT engine, it can be more efficient to export the engine node and run it using one of the other APIs. You can find instructions to do this in the [Torch-TRT documentation](https://pytorch.org/TensorRT/index.html).\n",
"\n",
"As an example, the Torch-TRT notebooks included with this guide use the Torch-TRT runtime.\n",
"\n",
"### TRITON Inference Server\n",
"\n",
"__Use this when:__\n",
"- You want to serve your models over HTTP or gRPC\n",
"- You want to load balance across multiple models or copies of models across GPUs to minimze latency and make better use of the GPU\n",
"- You want to have multiple models running efficiently on a single GPU at the same time\n",
"- You want to serve a variety of models converted using a variety of converters and frameworks (including Torch-TRT and ONNX) through a uniform interface\n",
"- You need serving support but are using PyTorch, another framework, or the ONNX path in general\n",
"\n",
"__More info:__\n",
"\n",
"\n",
"TRITON is an open source inference serving software that lets teams deploy trained AI models from any framework (TensorFlow, TensorRT, PyTorch, ONNX Runtime, or a custom framework), from local storage or Google Cloud Platform or AWS S3 on any GPU- or CPU-based infrastructure (cloud, data center, or edge). It is a flexible project with several unique features - such as concurrent model execution of both heterogeneous models and multiple copies of the same model (multiple model copies can reduce latency further) as well as load balancing and model analysis. It is a good option if you need to serve your models over HTTP - such as in a cloud inferencing solution.\n",
" \n",
"You can find the TRITON home page [here](https://developer.nvidia.com/triton-inference-server), and the documentation [here](https://docs.nvidia.com/deeplearning/triton-inference-server/user-guide/docs/)."
]
}
],
"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.10.12"
}
},
"nbformat": 4,
"nbformat_minor": 4
}
quickstart/IntroNotebooks/helper.py
+111
View File
@@ -0,0 +1,111 @@
#
# SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
from tensorflow.python.compiler.tensorrt import trt_convert as tf_trt
from tensorflow.python.saved_model import tag_constants
import tensorflow as tf
import tensorrt as trt
import numpy as np
precision_dict = {
"FP32": tf_trt.TrtPrecisionMode.FP32,
"FP16": tf_trt.TrtPrecisionMode.FP16,
"INT8": tf_trt.TrtPrecisionMode.INT8,
}
# For TF-TRT:
class OptimizedModel():
def __init__(self, saved_model_dir = None):
self.loaded_model_fn = None
if not saved_model_dir is None:
self.load_model(saved_model_dir)
def predict(self, input_data):
if self.loaded_model_fn is None:
raise(Exception("Haven't loaded a model"))
x = tf.constant(input_data.astype('float32'))
labeling = self.loaded_model_fn(x)
try:
preds = labeling['predictions'].numpy()
except:
try:
preds = labeling['probs'].numpy()
except:
try:
preds = labeling[next(iter(labeling.keys()))]
except:
raise(Exception("Failed to get predictions from saved model object"))
return preds
def load_model(self, saved_model_dir):
saved_model_loaded = tf.saved_model.load(saved_model_dir, tags=[tag_constants.SERVING])
wrapper_fp32 = saved_model_loaded.signatures['serving_default']
self.loaded_model_fn = wrapper_fp32
class ModelOptimizer():
def __init__(self, input_saved_model_dir, calibration_data=None):
self.input_saved_model_dir = input_saved_model_dir
self.calibration_data = None
self.loaded_model = None
if not calibration_data is None:
self.set_calibration_data(calibration_data)
def set_calibration_data(self, calibration_data):
def calibration_input_fn():
yield (tf.constant(calibration_data.astype('float32')), )
self.calibration_data = calibration_input_fn
def convert(self, output_saved_model_dir, precision="FP32", max_workspace_size_bytes=8000000000, **kwargs):
if precision == "INT8" and self.calibration_data is None:
raise(Exception("No calibration data set!"))
trt_precision = precision_dict[precision]
conversion_params = tf_trt.DEFAULT_TRT_CONVERSION_PARAMS._replace(precision_mode=trt_precision,
max_workspace_size_bytes=max_workspace_size_bytes,
use_calibration= precision == "INT8")
converter = tf_trt.TrtGraphConverterV2(input_saved_model_dir=self.input_saved_model_dir,
conversion_params=conversion_params)
if precision == "INT8":
converter.convert(calibration_input_fn=self.calibration_data)
else:
converter.convert()
converter.save(output_saved_model_dir=output_saved_model_dir)
return OptimizedModel(output_saved_model_dir)
def predict(self, input_data):
if self.loaded_model is None:
self.load_default_model()
return self.loaded_model.predict(input_data)
def load_default_model(self):
self.loaded_model = tf.keras.models.load_model('resnet50_saved_model')
quickstart/IntroNotebooks/images/onnx_onnx.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 110 KiB

quickstart/IntroNotebooks/images/pytorch_onnx.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 110 KiB

quickstart/IntroNotebooks/images/tensorrt_landscape.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 95 KiB

quickstart/IntroNotebooks/images/tensorrt_workflow.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 370 KiB

quickstart/IntroNotebooks/onnx_helper.py
+145
View File
@@ -0,0 +1,145 @@
#
# SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
import numpy as np
import tensorrt as trt
import weakref
from cuda.bindings import runtime as cudart
def _cleanup_cuda_resources(d_input, d_output, stream):
"""cleanup function for CUDA resources"""
if d_input is not None:
cudart.cudaFree(d_input)
if d_output is not None:
cudart.cudaFree(d_output)
if stream is not None:
cudart.cudaStreamDestroy(stream)
class ONNXClassifierWrapper:
def __init__(self, file, target_dtype=np.float32):
self.stream = None
self.d_input = None
self.d_output = None
self.target_dtype = target_dtype
self.num_classes = 1000
self.load(file)
self._finalizer = weakref.finalize(self, _cleanup_cuda_resources, self.d_input, self.d_output, self.stream)
def load(self, file):
with open(file, "rb") as f:
self.runtime = trt.Runtime(trt.Logger(trt.Logger.WARNING))
self.engine = self.runtime.deserialize_cuda_engine(f.read())
self.context = self.engine.create_execution_context()
def allocate_memory(self, batch):
self.output = np.empty(
self.num_classes, dtype=self.target_dtype
) # Need to set both input and output precisions to FP16 to fully enable FP16
# allocate device memory
err, self.d_input = cudart.cudaMalloc(batch.nbytes)
if err != cudart.cudaError_t.cudaSuccess:
raise RuntimeError(f"Failed to allocate input memory: {cudart.cudaGetErrorString(err)}")
err, self.d_output = cudart.cudaMalloc(self.output.nbytes)
if err != cudart.cudaError_t.cudaSuccess:
raise RuntimeError(f"Failed to allocate output memory: {cudart.cudaGetErrorString(err)}")
tensor_names = [self.engine.get_tensor_name(i) for i in range(self.engine.num_io_tensors)]
assert len(tensor_names) == 2
self.context.set_tensor_address(tensor_names[0], int(self.d_input))
self.context.set_tensor_address(tensor_names[1], int(self.d_output))
err, self.stream = cudart.cudaStreamCreate()
if err != cudart.cudaError_t.cudaSuccess:
raise RuntimeError(f"Failed to create stream: {cudart.cudaGetErrorString(err)}")
# update the finalizer with new resources
self._finalizer.detach()
self._finalizer = weakref.finalize(self, _cleanup_cuda_resources, self.d_input, self.d_output, self.stream)
def predict(self, batch): # result gets copied into output
if self.stream is None:
self.allocate_memory(batch)
# transfer input data to device
err = cudart.cudaMemcpyAsync(
self.d_input, batch.ctypes.data, batch.nbytes, cudart.cudaMemcpyKind.cudaMemcpyHostToDevice, self.stream
)
if err != cudart.cudaError_t.cudaSuccess:
raise RuntimeError(f"Failed to copy input to device: {cudart.cudaGetErrorString(err)}")
# execute model
self.context.execute_async_v3(self.stream)
# transfer predictions back
err = cudart.cudaMemcpyAsync(
self.output.ctypes.data,
self.d_output,
self.output.nbytes,
cudart.cudaMemcpyKind.cudaMemcpyDeviceToHost,
self.stream,
)
if err != cudart.cudaError_t.cudaSuccess:
raise RuntimeError(f"Failed to copy output from device: {cudart.cudaGetErrorString(err)}")
# synchronize threads
err = cudart.cudaStreamSynchronize(self.stream)
if err != cudart.cudaError_t.cudaSuccess:
raise RuntimeError(f"Failed to synchronize stream: {cudart.cudaGetErrorString(err)}")
return self.output
def cleanup(self):
"""free allocated CUDA memory and destroy stream"""
if hasattr(self, "_finalizer"):
self._finalizer()
self.d_input = None
self.d_output = None
self.stream = None
def convert_onnx_to_engine(onnx_filename, engine_filename=None, max_workspace_size=1 << 30, fp16_mode=True):
logger = trt.Logger(trt.Logger.WARNING)
with trt.Builder(logger) as builder, builder.create_network() as network, trt.OnnxParser(
network, logger
) as parser, builder.create_builder_config() as builder_config:
builder_config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, max_workspace_size)
if fp16_mode:
builder_config.set_flag(trt.BuilderFlag.FP16)
print("Parsing ONNX file.")
with open(onnx_filename, "rb") as model:
if not parser.parse(model.read()):
for error in range(parser.num_errors):
print(parser.get_error(error))
print("Building TensorRT engine. This may take a few minutes.")
serialized_engine = builder.build_serialized_network(network, builder_config)
if engine_filename:
with open(engine_filename, "wb") as f:
f.write(serialized_engine)
return serialized_engine, logger