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文件夹中找到包含延迟在内的原始输出。

设置

注意：torch2trt依赖于TensorRT Python API。在Jetson上，这包含在最新的JetPack中。对于桌面版，请遵循TensorRT安装指南。您也可以尝试在NGC PyTorch Docker容器中安装torch2trt，可选择桌面版或Jetson版。

步骤1 - 安装torch2trt Python库

要安装torch2trt Python库，请执行以下命令

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

步骤2（可选） - 安装torch2trt插件库

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

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

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

注意：torch2trt现在将插件作为独立库使用CMake编译。这使得编译的TensorRT引擎更具可移植性。如果需要，仍可以通过调用python setup.py install --plugins来安装依赖于PyTorch的已弃用插件。

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

要安装带有torch2trt.contrib下实验性社区贡献功能的torch2trt，如量化感知训练（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卡镜像