torch2trt下载torch2trt源代码下载

torch2trt

Python

v0.5.0

下载

TORCH2TRT

您正在使用什么模型或希望与Tensorrt一起使用？随时在这里加入讨论。

Torch2trt是使用张力python api的张力转换器的pytorch。转换器是

易于使用 - 带有单个功能调用torch2trt模块转换
易于扩展 - 在Python中写下您自己的图层转换器，然后将其注册为@tensorrt_converter

如果您发现问题，请告诉我们！

请注意，此转换器的张力 / pytorch覆盖率有限。我们主要创建它是为了轻松优化Jetbot项目中使用的模型。如果您发现转换器对其他型号有用，请告诉我们。

用法

以下是一些用法示例，有关更多信息，请查看笔记本。

转变

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

# create some regular pytorch model...
model = alexnet ( pretrained = True ). eval (). cuda ()

# create example data
x = torch . ones (( 1 , 3 , 224 , 224 )). cuda ()

# convert to TensorRT feeding sample data as input
model_trt = torch2trt ( model , [ x ])

执行

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

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

# check the output against 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中的吞吐量。您可以在基准文件夹中找到包括延迟的原始输出。

模型	纳米（pytorch）	纳米（张力）	Xavier（Pytorch）	Xavier（张力）
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取决于张力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的独立库。这使得汇编的张力引擎更便宜。如果需要，可以通过调用python setup.py install --plugins来安装不推荐的插件（取决于Pytorch）。

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

要在torch2trt.contrib下使用实验社区的TORCH2TRT，例如量化意识培训（QAT）（ requires 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函数的参数和返回语句以及正在构建的张力网络。原始pytorch函数的输入张量被修改为具有属性_trt ，它是pytorch张量的张力。转换函数使用此_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 )