使用torch2trt直接将超分辨率模型Real-ESRGAN转为TensorRT

2021-12-07 14:33:20 阅读：645 来源： 互联网

标签：Real torch TensorRT torch2trt cuda model trt

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

NVIDIA官网下载地址：https://developer.nvidia.com/zh-cn/tensorrt
建议先注册为developer。
选择最新的版本: TensorRT8 。最新版本包括更多的兼容的操作。
选择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
下载好后解压: tar -xzvf TensorRT-8.2.1.8.Linux.x86_64-gnu.cuda-11.4.cudnn8.2.tar.gz。
将TensorRT的lib可添加到系统变量里vim ~/.bashrc然后export LD_LIBRARY_PATH=/root/TensorRT-8.2.1.8/lib:$LD_LIBRARY_PATH，最后source ~/.bashrc。
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)))

标签：Real,torch,TensorRT,torch2trt,cuda,model,trt
来源： https://blog.csdn.net/qq_29598161/article/details/121765839

本站声明： 1. iCode9 技术分享网（下文简称本站）提供的所有内容，仅供技术学习、探讨和分享；
2. 关于本站的所有留言、评论、转载及引用，纯属内容发起人的个人观点，与本站观点和立场无关；
3. 关于本站的所有言论和文字，纯属内容发起人的个人观点，与本站观点和立场无关；
4. 本站文章均是网友提供，不完全保证技术分享内容的完整性、准确性、时效性、风险性和版权归属；如您发现该文章侵犯了您的权益，可联系我们第一时间进行删除；
5. 本站为非盈利性的个人网站，所有内容不会用来进行牟利，也不会利用任何形式的广告来间接获益，纯粹是为了广大技术爱好者提供技术内容和技术思想的分享性交流网站。