使用 TensorRT 加速深度学习推理

这是的更新版本 如何用 TensorRT 加速深度学习推理 . 此版本从 PyTorch 模型开始，而不是 ONNX 模型，将示例应用程序升级为使用 TensorRT 7 ，并将 ResNet-50 分类模型替换为 UNet ，这是一个分段模型。

图 1 。 TensorRT 徽标

NVIDIA TensorRT 是一个用于深度学习推理的 SDK 。 TensorRT 提供 api 和解析器来从所有主要的深度学习框架中导入经过训练的模型。然后生成可部署在数据中心、汽车和嵌入式环境中的优化运行时引擎。

这篇文章简单介绍了如何使用 TensorRT 。您将学习如何在 GPU 上部署深度学习应用程序，从而提高吞吐量并减少推理过程中的延迟。它使用 C ++示例来将您通过将 PyTorch 模型转换成 ONX 模型并将其导入 TensorRT ，应用优化，并为数据中心环境生成高性能运行时引擎。

TensorRT 支持 C ++和 Python ；如果使用了两者，工作流讨论可能是有用的。如果您喜欢使用 Python ，请参见 TensorRT 文档中的 使用 Python API 。

深度学习应用于广泛的应用，如自然语言处理、推荐系统、图像和视频分析。随着越来越多的应用程序在生产中使用深度学习，对准确性和性能的要求导致了模型复杂性和规模的强劲增长。

安全关键型应用（如汽车）对深度学习模型的吞吐量和延迟提出了严格的要求。对于包括推荐系统在内的一些消费者应用程序也是如此。

TensorRT 旨在帮助这些用例部署深度学习。通过对每一个主要框架的支持， TensorRT 通过强大的优化、降低精度的使用和高效的内存使用，帮助以低延迟处理大量数据。

为了继续这篇文章，您需要一台具有 CUDA -TensorRT 功能的计算机，或者一个安装了 GPUs 的云实例。在 Linux 上，最容易开始的地方是从 GPU 容器注册表下载带有 TensorRT 集成的 GPU – 加速的 PyTorch 集装箱 。

示例应用程序使用来自 Kaggle 的 脑 MRI 分割数据 的输入数据来执行推断。

简单 TensorRT 示例

以下是此示例应用程序的四个步骤：

1. 将预训练图像分割 PyTorch 模型转换为 ONNX 。
2. 将 ONNX 模型导入到 TensorRT 中。
3. 应用优化并生成引擎。
4. 对 GPU 执行推理。

导入 ONNX 模型包括从磁盘上保存的文件加载它，并将其从其本机框架或格式转换为 TensorRT 网络。 盒子 是一个表示深度学习模型的标准，使其能够在框架之间进行转换。

许多框架，如 Caffe2 、 Chainer 、 CNTK 、 paddlate 、 PyTorch 和 MXNet 都支持 ONNX 格式。接下来，根据输入模型、目标 GPU 平台和其他指定的配置参数，构建优化的 TensorRT 引擎。最后一步是向 TensorRT 引擎提供输入数据以执行推理。

应用程序在 TensorRT 中使用以下组件：

• ONX parser : 将一个经过 PyTorch 训练的模型转换成 ONNX 格式作为输入，并在 TensorRT 中填充一个网络对象。
• 建设者： 使用 TensorRT 中的网络并生成针对目标平台优化的引擎。
• 发动机： 获取输入数据，执行推断，并发出推断输出。
• 记录器： 与生成器和引擎关联，以在构建和推理阶段捕获错误、警告和其他信息。

将预训练图像分割 PyTorch 模型转换为 ONNX

从 NGC 注册表中的 PyTorch 集装箱 开始，预先安装框架和 CUDA 组件，准备就绪。成功安装 PyTorch 容器后，运行以下命令下载运行此示例应用程序所需的所有内容（示例代码、测试输入数据和引用输出）、更新依赖项，并使用提供的 makefile 编译应用程序。

>> sudo apt-get install libprotobuf-dev protobuf-compiler # protobuf is a prerequisite library
>> git clone --recursive https://github.com/onnx/onnx.git # Pull the ONNX repository from GitHub >> cd onnx
>> mkdir build && cd build >> cmake .. # Compile and install ONNX
>> make # Use the ‘-j’ option for parallel jobs, for example, ‘make -j $(nproc)’ >> make install >> cd ../..
>> git clone https://github.com/parallel-forall/code-samples.git
>> cd code-samples/posts/TensorRT-introduction
>> make clean && make # Compile the TensorRT C++ code
>> cd ..
>> wget https://developer.download.nvidia.com/devblogs/speeding-up-unet.7z // Get the ONNX model and test the data
>> tar xvf speeding-up-unet.7z # Unpack the model data into the unet folder
>> cd unet
>> python create_network.py #Inside the unet folder, it creates the unet.onnx file

将 PyTorch – 训练的 UNet 模型转换为 ONNX ，如下代码示例所示：

import torch
from torch.autograd import Variable
import torch.onnx as torch_onnx
import onnx
def main(): input_shape = (3, 256, 256) model_onnx_path = "unet.onnx" dummy_input = Variable(torch.randn(1, *input_shape)) model = torch.hub.load('mateuszbuda/brain-segmentation-pytorch', 'unet', in_channels=3, out_channels=1, init_features=32, pretra