Torch2TRT编译和使用踩坑

前言

Torch2TRT是英伟达提供的开源Pytorch到TensorRT模型的转化工具。相对于其他Pytorch模型转TensorRT的方式，我认为这是最简单和容易上手的方式。但是该工具并不成熟，在安装和使用过程中有一些坑。

遇到的问题

1. fatal error: xxxxxx.h: No such file or directory

例如:

fatal error: cuda_runtime_api.h: No such file or directory
fatal error: NvInfer.h No such file or directory

上面是找不到CUDA的include文件
下面是找不到TensorRT的include文件

首先需要确定确定是否已经安装CUDA和TensorRT，如果没有安装，则需要安装CUDA和TensorRT。这点相关教程多如牛毛，不再赘述。如果已经安装CUDA和TensorRT，可以在环境变量中指定include目录，如CUDA一般为/usr/local/cuda/include

export C_INCLUDE_PATH=$C_INCLUDE_PATH:/usr/local/cuda/include
export CPLUS_INCLUDE_PATH=$CPLUS_INCLUDE_PATH:/usr/local/cuda/include
export C_INCLUDE_PATH=$C_INCLUDE_PATH:<TesnsorRT include 目录>
export CPLUS_INCLUDE_PATH=$CPLUS_INCLUDE_PATH:<TesnsorRT include 目录>

另一种方法是修改项目的setup.py文件来指定include目录。如原文件在第9行函数def trt_inc_dir()中指定了TensorRT的include地址,我们需要修改内容指定自己电脑中TensorRT地址,例如:

def trt_inc_dir():return "/home/xxx/TensorRT-8.6.1.6/include"

如果缺少CUDA的include目录,可在原文件第29行添加，例如：

plugins_ext_module = CUDAExtension(name='plugins',sources=['torch2trt/plugins/plugins.cpp'],include_dirs=[trt_inc_dir(),'/usr/local/cuda/include'#在这里添加],library_dirs=[trt_lib_dir(),'/usr/local/cuda/lib64'],libraries=['nvinfer'],extra_compile_args={'cxx': compile_args_cxx,'nvcc': []})

2. /usr/bin/ld: cannot find -lxxxx: No such file or directory

例如：

/usr/bin/ld: cannot find -lcudart: No such file or directory
/usr/bin/ld: cannot find -lnvinfer: No such file or directory

上面是找不到CUDA的库文件
下面是找不到TensorRT的库文件

首先同样需要确定确定是否已经安装CUDA和TensorRT。如果已经安装，则需要修改项目的setup.py文件来指定lib目录。如果缺少TensorRT的库文件，则需要修改第12行的trt_lib_dir()函数指向自己的TensorRT的库目录，例如：

def trt_lib_dir():return "/home/xxx/TensorRT-8.6.1.6/lib"

如果缺少CUDA等其他库文件，可以在32行附近指定库文件目录，例如:

plugins_ext_module = CUDAExtension(name='plugins',sources=['torch2trt/plugins/plugins.cpp'],include_dirs=[trt_inc_dir(),'/usr/local/cuda/include'],library_dirs=[trt_lib_dir(),'/usr/local/cuda/lib64'#在这里添加],libraries=['nvinfer'],extra_compile_args={'cxx': compile_args_cxx,'nvcc': []})

3. module ‘collections’ has no attribute ‘Sequence’

与python3.10 不兼容。在Python3.10中，Sequence在collections.abc下，而不在collections下。
这一般是torch2trt/converters/interpolate.py文件报错，根据报错信息定位torch2trt/converters/interpolate.py文件地址，可修改该文件第5行：

#import collections
import collections.abc as collections

4. incompatible function arguments. The following argument types are supported

例如:

TypeError: (): incompatible function arguments. The following argument types are supported:1. (arg0: tensorrt.tensorrt.IConvolutionLayer, arg1: tensorrt.tensorrt.DimsHW) -> NoneInvoked with: <tensorrt.tensorrt.IConvolutionLayer object at 0x7f7aeaac30d8>, ([1, 1], [1, 1])

一般为torch2trt/converters/Conv2d.py文件中的bug，根据报错信息定位torch2trt/converters/Conv2d.py文件地址，可修改该文件第40行上下：

layer = ctx.network.add_convolution_nd(input=input_trt,num_output_maps=out_channels,kernel_shape=kernel_size,kernel=kernel,bias=bias)#  在这添加下面的8行代码print("Before stride"+str(stride))if isinstance(stride[0], int)==False:stride=(stride[0][0],stride[1][1])print("After stride"+str(stride))print('Before padding= '+str(padding))if isinstance(padding[0], int)==False:padding=(padding[0][0],padding[1][1])print('After padding= '+str(padding)layer.stride_nd = stridelayer.padding_nd = paddinglayer.dilation_nd = dilation