Environment

• GPU: Tesla T4
  • Driver Version: 460.91.03
• CUDA: 11.1
• CuDNN: 8.0.05
• PyTorch: 1.8.0
• Python: 3.8
• OS: ubuntu 18.04

如果安装的cuda是dev版本:
nvcc --version或者ls -all /usr/local/ 查看cuda软连接的cuda版本。
cat /usr/local/cuda-11.0/include/cudnn_version.h 查看cudnn版本。
使用torch查看使用的版本: torch.version.cuda 和 torch.backends.cudnn.version()

如果安装的是runtime版本，则不会有/usr/local/cuda目录，我没有测试这种情况下能否使用TensorRT。

上面是我测试用的环境，在你的环境中确保cuda、cudnn版本和gpu驱动匹配。

同时，确认系统变量正确设置:

vim ~/.bashrc

export PATH=/usr/local/cuda/bin:$PATH
export LD_LIBRARY_PATH=/usr/local/cuda/lib64:$LD_LIBRARY_PATH

source ~/.bashrc

一般使用NVIDIA官方的CUDA安装程序会自动设置，如果没有自己手动添加。

Installing TensorRT

0. NVIDIA官网下载地址：https://developer.nvidia.com/zh-cn/tensorrt
1. 建议先注册为developer。
2. 选择最新的版本: TensorRT8 。最新版本包括更多的兼容的操作。
3. 选择Tar包安装方式，直接简单，如: TensorRT 8.2 GA for Linux x86_64 and CUDA 11.0, 11.1, 11.2, 11.3, 11.4 and 11.5 TAR Package
4. 下载好后解压: tar -xzvf TensorRT-8.2.1.8.Linux.x86_64-gnu.cuda-11.4.cudnn8.2.tar.gz。
5. 将TensorRT的lib可添加到系统变量里vim ~/.bashrc然后export LD_LIBRARY_PATH=/root/TensorRT-8.2.1.8/lib:$LD_LIBRARY_PATH，最后source ~/.bashrc。
6. pip安装tensorrt: 到tensorrt解压目录下，pip install python/tensorrt-8.2.1.8-cp38-none-linux_x86_64.whl，用的python3.8，所以选cp38。
   此外还需要安装graphsurgeon，pip install graphsurgeon/graphsurgeon-0.4.5-py2.py3-none-any.whl。

Installing torch2trt

项目地址: https://github.com/NVIDIA-AI-IOT/torch2trt
文档地址: https://nvidia-ai-iot.github.io/torch2trt/v0.3.0/

由于在最新的torch2trt 0.3中已经包含torch.nn.functional.interpolate操作，所以选择无插件的安装方式

git clone https://github.com/NVIDIA-AI-IOT/torch2trt
cd torch2trt
python setup.py install

此外，这里也记录下插件的安装方式(当有些pytorch的操作在torch2trt未实现需要手动以插件方式实现):

git clone https://github.com/NVIDIA-AI-IOT/torch2trt
cd torch2trt
python setup.py install --plugins

报错: NvInfer.h: No such file or directory
解决方法:
编辑setup.py文件

include_dirs=[
    trt_inc_dir(),
    'your/path/TensorRT-8.2.1.8/include'       # add include directories
],
library_dirs=[
    trt_lib_dir(),
    'your/path/TensorRT-8.2.1.8/lib'           # add link directories
],

报错: error: invalid new-expression of abstract class type ‘torch2trt::GroupNormPlugin’
暂无解决方法，可能是TensorRT版本与torch2trt不匹配？还未测试。

Testing Basic Usage

我们使用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 example data
x = torch.ones((1, 3, 224, 224)).cuda()

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

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

# check the output against PyTorch
print(torch.max(torch.abs(y - y_trt)))

如果未报错则说明配置成功，测试输出结果:

tensor(1.0729e-06, device='cuda:0', grad_fn=<MaxBackward1>)

Testing Conversion of RRDBNet

import time
import torch
import os

from basicsr.archs.rrdbnet_arch import RRDBNet
from super_resolution.real_esrgan.utils import RealESRGANer
from torch2trt import torch2trt


# create some regular pytorch model...
def get_sr_model(tile_size=None, model_path_prefix='', scale=2):
    rrdb_model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=scale)
    model = RealESRGANer(scale=scale, model_path=os.path.join(model_path_prefix,
                                                              'src/pretrained_models/real_esrgan/RealESRGAN_x2plus.pth'),
                         model=rrdb_model, tile_size=tile_size, tile_pad=10, pre_pad=10)
    return model


model = get_sr_model(tile_size=None, model_path_prefix='../', scale=2).model

# create example data
x = torch.randn((1, 3, 256, 256)).cuda().half()

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

torch.save(model_trt.state_dict(), 'RealESRGAN_x2plus_trt.pth')

s = time.time()
y = model(x)
print(time.time() - s)

s = time.time()
y_trt = model_trt(x)
print(time.time() - s)
# check the output against PyTorch
print(torch.max(torch.abs(y - y_trt)))