torch2trt

您正在使用或希望使用哪些模型与 TensorRT 配合？欢迎加入讨论 这里。

torch2trt 是一个基于 TensorRT Python API 的 PyTorch 到 TensorRT 转换器。该转换器具有以下特点：

• 易于使用 - 仅需调用单个函数 torch2trt 即可转换模块。

• 易于扩展 - 您可以用 Python 编写自己的层转换器，并使用 @tensorrt_converter 注册它。

如果您发现任何问题，请 告知我们！

请注意，此转换器对 TensorRT 和 PyTorch 的支持范围有限。我们主要创建它是为了方便优化 JetBot 项目中使用的模型。如果您发现该转换器对其他模型也有帮助，请 告知我们。

使用方法

以下是一些使用示例，更多内容请参阅 notebooks。

转换

import torch
from torch2trt import torch2trt
from torchvision.models.alexnet import alexnet

# 创建一个普通的 PyTorch 模型...
model = alexnet(pretrained=True).eval().cuda()

# 创建示例数据
x = torch.ones((1, 3, 224, 224)).cuda()

# 使用示例数据作为输入将模型转换为 TensorRT
model_trt = torch2trt(model, [x])

执行

我们可以像执行原始 PyTorch 模型一样执行返回的 TRTModule。

y = model(x)
y_trt = model_trt(x)

# 将输出与 PyTorch 进行比较
print(torch.max(torch.abs(y - y_trt)))

保存和加载

我们可以将模型保存为 state_dict。

torch.save(model_trt.state_dict(), 'alexnet_trt.pth')

然后可以将保存的模型加载到 TRTModule 中。

from torch2trt import TRTModule

model_trt = TRTModule()
model_trt.load_state_dict(torch.load('alexnet_trt.pth'))

模型

我们使用 test.sh 脚本对这些模型进行了测试。您可以通过运行以下命令生成结果：

./test.sh TEST_OUTPUT.md

下面的结果显示了以 FPS 为单位的吞吐量。包含延迟在内的原始输出可以在 benchmarks 文件夹 中找到。

模型	Nano (PyTorch)	Nano (TensorRT)	Xavier (PyTorch)	Xavier (TensorRT)
alexnet	46.4	69.9	250	580
squeezenet1_0	44	137	130	890
squeezenet1_1	76.6	248	132	1390
resnet18	29.4	90.2	140	712
resnet34	15.5	50.7	79.2	393
resnet50	12.4	34.2	55.5	312
resnet101	7.18	19.9	28.5	170
resnet152	4.96	14.1	18.9	121
densenet121	11.5	41.9	23.0	168
densenet169	8.25	33.2	16.3	118
densenet201	6.84	25.4	13.3	90.9
densenet161	4.71	15.6	17.2	82.4
vgg11	8.9	18.3	85.2	201
vgg13	6.53	14.7	71.9	166
vgg16	5.09	11.9	61.7	139
vgg19			54.1	121
vgg11_bn	8.74	18.4	81.8	201
vgg13_bn	6.31	14.8	68.0	166
vgg16_bn	4.96	12.0	58.5	140
vgg19_bn			51.4	121

安装

注意：torch2trt 依赖于 TensorRT 的 Python API。在 Jetson 上，这已包含在最新的 JetPack 中。对于桌面端，请按照 TensorRT 安装指南 进行安装。您也可以尝试在 NGC 的 PyTorch Docker 容器中安装 torch2trt，适用于 桌面 或 Jetson。

第一步 - 安装 torch2trt Python 库

要安装 torch2trt Python 库，请执行以下操作：

git clone https://github.com/NVIDIA-AI-IOT/torch2trt
cd torch2trt
python setup.py install

第二步（可选）- 安装 torch2trt 插件库

要安装 torch2trt 插件库，请执行以下命令：

cmake -B build . && cmake --build build --target install && ldconfig

这包括对某些可能未被 TensorRT 原生支持的层的支持。一旦系统中找到此库，torch2trt 中相关的层转换器就会自动启用。

注意：现在 torch2trt 将插件维护为一个独立的库，并使用 CMake 进行编译。这使得编译后的 TensorRT 引擎更具可移植性。如果需要，仍然可以通过运行 python setup.py install --plugins 来安装已弃用的插件（这些插件依赖于 PyTorch）。

第三步（可选）- 安装实验性的社区贡献功能

要安装带有实验性社区贡献功能的 torch2trt（位于 torch2trt.contrib 中），例如量化感知训练 (QAT)（需要 TensorRT>=7.0），请执行以下操作：

git clone https://github.com/NVIDIA-AI-IOT/torch2trt
cd torch2trt/scripts    
bash build_contrib.sh   

这将使您能够运行位于 此处 的 QAT 示例。

工作原理

该转换器通过将转换函数（如 convert_ReLU）附加到原始的 PyTorch 函数调用（如 torch.nn.ReLU.forward）来工作。示例输入数据会像以前一样通过网络传递，但每当遇到已注册的函数（torch.nn.ReLU.forward）时，相应的转换器（convert_ReLU）也会被调用。转换器会接收原始 PyTorch 函数的参数和返回值，以及正在构建的 TensorRT 网络。原始 PyTorch 函数的输入张量会被修改，添加一个 _trt 属性，该属性是 PyTorch 张量对应的 TensorRT 版本。转换函数会利用这个 _trt 属性向 TensorRT 网络中添加层，并为相关的输出张量设置 _trt 属性。当模型完全执行完毕后，最终的输出张量会被标记为 TensorRT 网络的输出，随后构建优化后的 TensorRT 引擎。

如何添加（或覆盖）转换器

下面我们展示如何使用 TensorRT 的 Python API 为 ReLU 模块添加一个转换器。

import tensorrt as trt
from torch2trt import tensorrt_converter

@tensorrt_converter('torch.nn.ReLU.forward')
def convert_ReLU(ctx):
    input = ctx.method_args[1]
    output = ctx.method_return
    layer = ctx.network.add_activation(input=input._trt, type=trt.ActivationType.RELU)  
    output._trt = layer.get_output(0)

该转换器接受一个参数，即 ConversionContext 对象，其中包含以下内容：

• ctx.network - 当前正在构建的 TensorRT 网络。
• ctx.method_args - 传递给指定 PyTorch 函数的位置参数。相关输入张量会设置 _trt 属性。
• ctx.method_kwargs - 传递给指定 PyTorch 函数的关键字参数。
• ctx.method_return - 指定 PyTorch 函数返回的值。转换器需要在适当的地方设置 _trt 属性。

更多示例请参阅 此文件夹。

参阅

• JetBot - 基于 NVIDIA Jetson Nano 的教育型 AI 机器人
• JetRacer - 使用 NVIDIA Jetson Nano 的教育型 AI 赛车
• JetCam - 面向 NVIDIA Jetson 的易用 Python 相机接口
• JetCard - 用于基于 NVIDIA Jetson Nano 的 Web 编程 AI 项目的 SD 卡镜像

torch2trt 快速上手指南

torch2trt 是一个基于 Python API 的 PyTorch 到 TensorRT 转换工具。它允许开发者通过单个函数调用轻松将 PyTorch 模型转换为优化的 TensorRT 引擎，显著提升在 NVIDIA Jetson 或桌面 GPU 上的推理速度。

环境准备

在使用 torch2trt 之前，请确保满足以下系统和依赖要求：

• 操作系统: Linux (推荐 Ubuntu)
• 核心依赖:
  • PyTorch: 已安装并配置好 CUDA 支持。
  • TensorRT: 必须安装 TensorRT 及其 Python API。
    • Jetson 平台: 包含在最新的 JetPack SDK 中，无需单独安装。
    • 桌面平台 (x86_64): 请参考 NVIDIA TensorRT 安装指南 进行安装。
• 编译工具 (可选，用于安装插件): cmake, g++, make。

提示: 对于桌面用户，也可以直接使用 NVIDIA NGC 提供的 PyTorch Docker 容器（包含预装的环境）：

• Desktop PyTorch Container
• Jetson PyTorch Container

安装步骤

1. 安装 torch2trt 核心库

克隆仓库并运行安装脚本：

git clone https://github.com/NVIDIA-AI-IOT/torch2trt
cd torch2trt
python setup.py install

2. (可选) 安装插件库

如果需要支持某些 TensorRT 原生不支持的层，建议编译安装 C++ 插件库。这将使生成的引擎更具可移植性。

cmake -B build . && cmake --build build --target install && ldconfig

3. (可选) 安装社区贡献的实验性功能

如需使用量化感知训练 (QAT) 等实验性功能（需 TensorRT >= 7.0）：

cd scripts    
bash build_contrib.sh   

基本使用

以下是将 PyTorch 模型转换为 TensorRT 模型并进行推理的最简流程。

1. 模型转换

导入必要的库，加载预训练的 PyTorch 模型，准备示例输入数据，然后调用 torch2trt 进行转换。

import torch
from torch2trt import torch2trt
from torchvision.models.alexnet import alexnet

# 创建普通的 PyTorch 模型 (例如 AlexNet)
model = alexnet(pretrained=True).eval().cuda()

# 创建示例输入数据 (形状需与实际推理一致)
x = torch.ones((1, 3, 224, 224)).cuda()

# 执行转换：传入模型和示例数据
model_trt = torch2trt(model, [x])

2. 执行推理

转换后的 model_trt 是一个 TRTModule，其调用方式与原 PyTorch 模型完全一致。

# 原始 PyTorch 推理
y = model(x)

# TensorRT 加速推理
y_trt = model_trt(x)

# (可选) 验证输出差异
print(torch.max(torch.abs(y - y_trt)))

3. 保存与加载

可以将优化后的 TensorRT 引擎保存为文件，以便后续直接加载使用，无需重新转换。

保存模型：

torch.save(model_trt.state_dict(), 'alexnet_trt.pth')

加载模型：

from torch2trt import TRTModule

model_trt = TRTModule()
model_trt.load_state_dict(torch.load('alexnet_trt.pth'))

# 加载后即可直接使用
# y = model_trt(x)