TensorRT安装

按照官方文档来，一般没啥问题。我使用的是tar格式安装。

由于TensorRT 7对CUDA版本的要求，我原先服务器的CUDA版本为10.1，因此新安装了CUDA11.0并做了切换（可以参考这里）
TensorRT需要安装CuDNN（同样参照CuDNN官方文档）。
20210416：需要提醒的是，文档上给了tar文件安装指引，但是CuDNN下载页只提供了后缀为.solitairetheme8和.deb的文件下载。如果使用tar文件方式，需要下载.solitairetheme8格式的那个，下载后将其后缀改为.tgz后按tar文件安装方式来即可。

遇到的问题

Q：ImportError: libnvinfer.so.7: cannot open shared object file: No such file or directory

A：参考，这是由于tensorrt的路径没有添加到自己的环境变量。解决方法为：

  
# 打开~/.bashrc
vim ~/.bashrc

# 点击insert进入编辑模式，将下面一行加入（注意替换自己的tensorRT路径）
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:path/to/your_tensorRT/lib
# 编辑完成后，依次键入Esc, :wq退出

# 刷新
source ~/.bashrc     # 这一步可能要等较长时间

另，关于Linux的环境变量，可参考这篇

Pytorch 模型向TensorRT部署

一个很好的TensorRT入门资料

目前已经尝试了两种可行的方案：

使用NVIDIA-AI-IOT/torch2trt
pytorch \(\rightarrow\) onnx \(\rightarrow\) TensorRT

使用torch2trt

github

document

假设现在已经有了模型的定义文件model.py和pytorch权重文件weight.pt

基础操作

torch2trt基础操作使用起来非常简单

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

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

# create a dummy input data
x = torch.ones((1, 3, 224, 224)).cuda()

# convert pytorch model to trt model (fp32) 
model_trt = torch2trt(model, [x])

# excute and check the output of the converted trt_model
y = model(x)
y_trt = model_trt(x)
print(torch.max(torch.abs(y - y_trt)))

# save the trt model as a state_dict.
torch.save(model_trt.state_dict(), 'alexnet_trt.pth')

# load the saved model into a TRTModule
from torch2trt import TRTModule

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

一些更高级一些的convert操作

参照<torch2trt_root>/torch2trt/torch2trt.py#L482

  
def torch2trt(module, 
              inputs, 
              input_names=None, 
              output_names=None, 
              log_level=trt.Logger.ERROR, 
              max_batch_size=1,
              fp16_mode=False, 
              max_workspace_size=1<<25, 
              strict_type_constraints=False, 
              keep_network=True, 
              int8_mode=False, 
              int8_calib_dataset=None,
              int8_calib_algorithm=DEFAULT_CALIBRATION_ALGORITHM,
              int8_calib_batch_size=1,
              use_onnx=False,
              **kwargs):
    ...

设定max_batch_size

  
# convert pytorch model to trt model (fp32) 
Max_batch_size = 8
model_trt = torch2trt(model, [x], max_batch_size=Max_batch_size)

convert to fp16 trt model

  
# convert pytorch model to trt model (fp16) 
model_trt = torch2trt(model, [x], fp16_mode=True)

convert to int8 model

转化为int8模型需要构建calibrate数据集进行校正，按照这里使用Pytorch的dataset类构建一个cali_dataset传入即可。

  
cali_dataset = Your_defined_cali_dataset
# convert pytorch model to trt model (int8)
model_trt = torch2trt(model, [x], 
                      int8_mode=True,
                      int8_int8_calib_dataset=cali_dataset)

注：官方的__getitem__(self, index)实现为

  
def __getitem__(self, idx):
    image, _ = self.dataset[idx]
    image = image[None, ...]  # add batch dimension
    return [image]

但我运行起来会有维度不匹配的问题，我在实际使用时使用了如下定义

  
def __getitem__(self, idx):
    image, _ = self.dataset[idx]
    if self.transform is not None:
        image = self.transform(image)
        
    return [image]

扩展converter

目前torch2trt支持的所有converter可以看这里，但并没有涵盖所有TensorRT已经支持的所有算子。当所转化的模型中包含目前torch2trt的converter没有实现的算子，则会报错并模型转化失败，比如对torch.nn.MaxPool3d会出现

warning: encountered known unsupported method torch.max_pool3d
warning: encountered known unsupported method torch.nn.fucntional.max_pool3d
[TensorRT] ERROR: INVALID_ARGUMENT: Cannot find binding of given name: input_0

理论上，只要是TensorRT已经支持的算子，即使是torch2trt没有定义对应的converter，我们参照torch2trt已经实现的converer，自己也可以很方便地将其实现，官方给了一个ReLU算子的converter例子。

自定义一个converter的流程为：

在<torch2trt_root>/torch2trt/conveters/下新建一个<your_custom_converter>.py
在<your_custom_converter>.py中实现converter
在<torch2trt_root>/torch2trt/conveters/__init__.py中import新定义的converter
重新安装torch2trt：python setup.py install

以下记录下MaxPool3d的converter实现流程，非常方便，很大程度上参照了avg_pool.py中avg_pool3d的converter的实现。

在<torch2trt_root>/torch2trt/conveters/下新建一个max_pool3d.py

在max_pool3d.py实现torch.nn.MaxPool3d的converter

  
# A CUSTOM ONE #
from torch2trt.torch2trt import *
from torch2trt.module_test import add_module_test
   
@tensorrt_converter('torch.nn.functional.max_pool3d', enabled=trt_version() >= '7.0')
def convert_max_pool_trt7(ctx):
    # print(ctx.method_args)
    # parse args
    # refer to https://pytorch.org/docs/stable/generated/torch.nn.MaxPool3d.html
    input = get_arg(ctx, 'input', pos=0, default=None)
    kernel_size = get_arg(ctx, 'kernel_size', pos=1, default=None)
    stride = get_arg(ctx, 'stride', pos=2, default=None)
    padding = get_arg(ctx, 'padding', pos=3, default=0)
    dilation = get_arg(ctx, 'dilation', pos=4, default=False)
    return_indices = get_arg(ctx, 'return_indices', pos=5, default=True)
    ceil_mode = get_arg(ctx, 'ceil_mode', pos=6, default=True)
    # print(kernel_size, stride, padding, ceil_mode, count_include_pad)
   
    # get input trt tensor (or create constant if it doesn't exist)
    input_trt = add_missing_trt_tensors(ctx.network, [input])[0]
    output = ctx.method_return
   
    input_dim = input.dim() - 2
   
    # get kernel size
    if not isinstance(kernel_size, tuple):
        kernel_size = (kernel_size,) * input_dim
   
    # get stride
    if not isinstance(stride, tuple):
        stride = (stride,) * input_dim
   
    # get padding
    if not isinstance(padding, tuple):
        padding = (padding,) * input_dim
   	
    # For type, check the TensorRT Python API document
    layer = ctx.network.add_pooling_nd(
        input=input_trt, type=trt.PoolingType.MAX, window_size=kernel_size)
   
    layer.stride_nd = stride
    layer.padding_nd = padding
       
    if ceil_mode:
        layer.padding_mode = trt.PaddingMode.EXPLICIT_ROUND_UP
   
    output._trt = layer.get_output(0)

看起来很简单的对不？需要提醒的是最开始获取当前layer的参数

  
input = get_arg(ctx, 'input', pos=0, default=None)
...
ceil_mode = get_arg(ctx, 'ceil_mode', pos=6, default=True)

从torch2trt.torch2trt.py看到get_arg的定义

  
def get_arg(ctx, name, pos, default):
    if name in ctx.method_kwargs:
        return ctx.method_kwargs[name]
    elif len(ctx.method_args) > pos:
        return ctx.method_args[pos]
    else:
        return default

这里的method_kwargs和Pytorch对应的算子输入参数是一一对应的！（从这个函数的pos参数来看，除了第0位挤入了input，其他参数都是依次向后排开的）所以要转化某一个算子，需要到Pytorch官网去看看这个算子的参数表，比如torch.nn.MaxPool3d

在<torch2trt_root>/torch2trt/conveters/__init__.py中import新定义的converter
1 from .max_pool3d import *
重新安装torch2trt：python setup.py install

再次尝试转化含torch.nn.MaxPool3d的模型，报错消失，转化后的模型运行也正常，问题解决~！

Pytorch \(\rightarrow\) ONNX \(\rightarrow\) TensorRT

其他

竟然还可以用pip安装tensorRT?! here

TensorRT使用