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) -> None

Invoked 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 = stride
    layer.padding_nd = padding
    layer.dilation_nd = dilation