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版。

步骤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)