使用TensorRT加速超分辨率模型Real-ESRGAN

最新推荐文章于 2024-04-12 09:11:12 发布
最新推荐文章于 2024-04-12 09:11:12 发布
阅读量3k 收藏 24
点赞数 3
文章标签： pytorch 深度学习 神经网络
版权声明：本文为博主原创文章，遵循 CC 4.0 BY-SA 版权协议，转载请附上原文出处链接和本声明。
本文链接：https://blog.csdn.net/qq_29598161/article/details/121765839
版权

Introduction

TensorRT是NVIDIA提供在其GPU上进行推理加速的工具。
这里以超分辨率模型 Real-ESRGAN 为例介绍3种方式将PyTorch模型转为TensorRT优化的模型并测试性能。

  • torch2trt 是NVIDIA-AI-IOT开源的项目，旨在提供便利的API将PyTorch模型转为TensorRT模型，
    但是目前不支持动态尺寸的输入。
  • Torch-TensorRT 是NVIDIA开源的项目，也是一个直接将将PyTorch模型转为TensorRT模型的工具,
    它的扩展性会比torch2trt更好，主要原理是将PyTorch转为Torch-TensorRT支持的TorchScript。我建议使用这种方式，安装简单并且支持动态输入尺寸。
  • onnx-tensorrt 是ONNX开源的项目，将onnx模型转为TensorRT模型。

Environment

  • GPU: Tesla T4
    • Driver Version: 470.82.01
  • CUDA: 11.3
  • cuDNN: 8.2.0
  • 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/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-11.3/bin:$PATH
export LD_LIBRARY_PATH=/usr/local/cuda-11.3/lib64:$LD_LIBRARY_PATH

source ~/.bashrc

Installing TensorRT

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

PyTorch Convert TensorRT

Real-ESRGAN推理时主要使用RRDB-Net(basicsr/archs/rrdbnet_arch.py)，当scale=2时(本文也以scale=2为例)，其中有个pixel_unshuffle方法会导致转onnx有问题，
如这个issue，
我这里的解决方法是直接将pixel_unshuffle移出forward，事实上pixel_unshuffle的作用只是将分辨率压缩到到通道维度上，
如[b, 3, h, w] -> [b, 3*4, h/2, w/2]。
修改后的RRDB-Net(新建一个文件test/rrdb_net.py并引用):

import torch
from torch import nn as nn
from torch.nn import functional as F
from basicsr.archs.arch_util import default_init_weights, make_layer

class ResidualDenseBlock(nn.Module):
    def __init__(self, num_feat=64, num_grow_ch=32):
        super(ResidualDenseBlock, self).__init__()
        self.conv1 = nn.Conv2d(num_feat, num_grow_ch, 3, 1, 1)
        self.conv2 = nn.Conv2d(num_feat + num_grow_ch, num_grow_ch, 3, 1, 1)
        self.conv3 = nn.Conv2d(num_feat + 2 * num_grow_ch, num_grow_ch, 3, 1, 1)
        self.conv4 = nn.Conv2d(num_feat + 3 * num_grow_ch, num_grow_ch, 3, 1, 1)
        self.conv5 = nn.Conv2d(num_feat + 4 * num_grow_ch, num_feat, 3, 1, 1)

        self.lrelu = nn.LeakyReLU(negative_slope=0.2, inplace=True)
        # initialization
        default_init_weights([self.conv1, self.conv2, self.conv3, self.conv4, self.conv5], 0.1)

    def forward(self, x):
        x1 = self.lrelu(self.conv1(x))
        x2 = self.lrelu(self.conv2(torch.cat((x, x1), 1)))
        x3 = self.lrelu(self.conv3(torch.cat((x, x1, x2), 1)))
        x4 = self.lrelu(self.conv4(torch.cat((x, x1, x2, x3), 1)))
        x5 = self.conv5(torch.cat((x, x1, x2, x3, x4), 1))
        # Emperically, we use 0.2 to scale the residual for better performance
        return x5 * 0.2 + x


class RRDB(nn.Module):
    def __init__(self, num_feat, num_grow_ch=32):
        super(RRDB, self).__init__()
        self.rdb1 = ResidualDenseBlock(num_feat, num_grow_ch)
        self.rdb2 = ResidualDenseBlock(num_feat, num_grow_ch)
        self.rdb3 = ResidualDenseBlock(num_feat, num_grow_ch)

    def forward(self, x):
        out = self.rdb1(x)
        out = self.rdb2(out)
        out = self.rdb3(out)
        # Emperically, we use 0.2 to scale the residual for better performance
        return out * 0.2 + x


class RRDBNet(nn.Module):
    def __init__(self, num_in_ch, num_out_ch, num_feat=64, num_block=23, num_grow_ch=32):
        super(RRDBNet, self).__init__()
        num_in_ch = num_in_ch * 4
        self.conv_first = nn.Conv2d(num_in_ch, num_feat, 3, 1, 1)
        self.body = make_layer(RRDB, num_block, num_feat=num_feat, num_grow_ch=num_grow_ch)
        self.conv_body = nn.Conv2d(num_feat, num_feat, 3, 1, 1)

        self.conv_up1 = nn.Conv2d(num_feat, num_feat, 3, 1, 1)
        self.conv_up2 = nn.Conv2d(num_feat, num_feat, 3, 1, 1)
        self.conv_hr = nn.Conv2d(num_feat, num_feat, 3, 1, 1)
        self.conv_last = nn.Conv2d(num_feat, num_out_ch, 3, 1, 1)

        self.lrelu = nn.LeakyReLU(negative_slope=0.2, inplace=True)

    def forward(self, x):
        feat = self.conv_first(x)
        body_feat = self.conv_body(self.body(feat))
        feat = feat + body_feat

        feat = self.lrelu(self.conv_up1(F.interpolate(feat, scale_factor=2, mode='nearest')))
        feat = self.lrelu(self.conv_up2(F.interpolate(feat, scale_factor=2, mode='nearest')))
        out = self.conv_last(self.lrelu(self.conv_hr(feat)))
        return out

因此需要在输入网络前需要先做这个pixel_unshuffle操作在prepare中，另外提供一些共用方法，例如benchmark测试和显示图片等(新建一个文件test/common.py):

import time
import torch
import cv2
import numpy as np

from test.rrdb_net import RRDBNet
import matplotlib.pyplot as plt


def get_sr_model():
    model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32)
    load_net = torch.load('src/pretrained_models/real_esrgan/RealESRGAN_x2plus.pth')
    model.load_state_dict(load_net['params_ema'], strict=True)
    model.eval()
    model = model.to('cuda').half()
    return model


def show(x: torch.Tensor):
    print(x.shape)
    x = x[0].clip(0, 1).cpu().detach().numpy().transpose(1, 2, 0).astype(np.float32)
    plt.imshow(x)
    plt.show()


@torch.no_grad()
def benchmark(model, x, warm_up=2, runs=10):
    print("Warm up ...")
    with torch.no_grad():
        for _ in range(warm_up):
            features = model(x)
    torch.cuda.synchronize()
    print("Start timing ...")
    timings = []
    with torch.no_grad():
        for i in range(1, runs + 1):
            start_time = time.time()
            features = model(x)
            torch.cuda.synchronize()
            end_time = time.time()
            timings.append(end_time - start_time)
            if i % 10 == 0:
                print('Iteration %d/%d, ave batch time %.2f ms' % (i, runs, np.mean(timings[i-10: i]) * 1000))

    print("Input shape:", x.shape)
    print("Output features size:", features.shape)
    print('Average batch time: %.2f ms' % (np.mean(timings) * 1000))


def test(x, model, name):
    torch.cuda.empty_cache()
    s = time.time()
    with torch.no_grad():
        y = model(x)
    torch.cuda.synchronize()
    print(name, time.time() - s)
    show(y)
    return y


def prepare(side=512):
    model = get_sr_model().eval()

    # create example data
    x = cv2.resize(cv2.imread('../src/inputs/sr/0014.jpeg'),
                   (side, side))[..., ::-1].transpose(2, 0, 1) / 255.0
    x = torch.from_numpy(x).cuda().half().unsqueeze(0)
    show(x)

    b, c, h, w = x.size()
    h = h//2
    w = w//2
    x = x.view(b, c, h, 2, w, 2).permute(0, 1, 3, 5, 2, 4).reshape(b, 12, h, w)
    return model, x

下面说明具体工具使用：

torch2trt

torch2trt 目前只能固定一种尺寸的输入，可以参考这个
torch2trt_dynamic 实现多尺寸。

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 of torch2trt

使用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

原始模型使用FP16精度，TensorRT也使用FP16

import torch
from torch2trt import torch2trt
from test.common import prepare, test

side = 1024
model, x = prepare(side)

torch.cuda.empty_cache()
with torch.no_grad():
    print('converting trt...')
    model_trt = torch2trt(model, [x], fp16_mode=True)

torch.save(model_trt.state_dict(), f't2trt_fp16_{side}.trt')

y = test(x, model, "model")
y_trt = test(x, model_trt, "model_trt")
print("error", torch.max(torch.abs(y - y_trt)))

结果：

torch.Size([1, 3, 512, 512])
Converting trt...
Converted! used time 239.82s
Testing model used 0.38s
torch.Size([1, 3, 1024, 1024])
Testing model_trt used 0.20s
torch.Size([1, 3, 1024, 1024])
Max error tensor(0.1553, device='cuda:0', dtype=torch.float16)

批量测试优化效果:

import torch

from torch2trt import TRTModule
from test.common import prepare, test, benchmark

side = 512
model, x = prepare(side)


torch.cuda.empty_cache()

model_trt = TRTModule()
model_trt.load_state_dict(torch.load(f'src/torch2trt_fp16_{side}.trt'))

y = test(x, model, "model")
y_trt = test(x, model_trt, "model_trt")
print("Max error", torch.max(torch.abs(y - y_trt)))

print('*'*100)
torch.cuda.empty_cache()
benchmark(model, x, warm_up=5, runs=50)

print('*'*100)
torch.cuda.empty_cache()
benchmark(model_trt, x, warm_up=5, runs=50)

结果：

torch.Size([1, 3, 512, 512])
[12/10/2021-15:06:46] [TRT] [W] TensorRT was linked against cuBLAS/cuBLAS LT 11.6.3 but loaded cuBLAS/cuBLAS LT 11.5.1
[12/10/2021-15:06:47] [TRT] [W] TensorRT was linked against cuDNN 8.2.1 but loaded cuDNN 8.2.0
[12/10/2021-15:06:47] [TRT] [W] TensorRT was linked against cuBLAS/cuBLAS LT 11.6.3 but loaded cuBLAS/cuBLAS LT 11.5.1
[12/10/2021-15:06:47] [TRT] [W] TensorRT was linked against cuDNN 8.2.1 but loaded cuDNN 8.2.0
Testing model used 0.40s
torch.Size([1, 3, 1024, 1024])
Testing model_trt used 0.19s
torch.Size([1, 3, 1024, 1024])
Max error tensor(0.1553, device='cuda:0', dtype=torch.float16)
****************************************************************************************************
Warm up ...
Start timing ...
Iteration 10/50, ave batch time 314.55 ms
Iteration 20/50, ave batch time 314.76 ms
Iteration 30/50, ave batch time 315.70 ms
Iteration 40/50, ave batch time 316.29 ms
Iteration 50/50, ave batch time 317.29 ms
Input shape: torch.Size([1, 12, 256, 256])
Output features size: torch.Size([1, 3, 1024, 1024])
Average batch time: 315.72 ms
****************************************************************************************************
Warm up ...
Start timing ...
Iteration 10/50, ave batch time 200.73 ms
Iteration 20/50, ave batch time 201.76 ms
Iteration 30/50, ave batch time 202.16 ms
Iteration 40/50, ave batch time 202.45 ms
Iteration 50/50, ave batch time 200.82 ms
Input shape: torch.Size([1, 12, 256, 256])
Output features size: torch.Size([1, 3, 1024, 1024])
Average batch time: 201.59 ms

可以看到使用TensorRT后时间减少了1/3左右。另外在3090上测试时间可以减少一半。

Torch-TensorRT

Installing Torch-TensorRT

python直接使用pip安装Torch-TensorRT 。

pip install torch-tensorrt -f https://github.com/NVIDIA/Torch-TensorRT/releases

Torch-TensorRT 的文档比较详细，其中包括c++部署。Torch-TensorRT 支持动态输入尺寸，需要设置min_shape、opt_shape、max_shape。

Testing Conversion of RRDBNet

Fixed input shape
import time
import torch_tensorrt
import torch
from test.common import prepare, test

side = 512
model, x = prepare(side)

compile_settings = {
    "inputs": [
        torch_tensorrt.Input(
            (1, 12, side // 2, side // 2),
            dtype=torch.half)
    ],
    "truncate_long_and_double": True,
    "enabled_precisions": {torch.half}  # Run with FP16
}

with torch.no_grad():
    torch.cuda.empty_cache()
    traced_model = torch.jit.trace(model, x)
    torch.cuda.empty_cache()
    print('Converting trt...')
    s = time.time()
    model_trt = torch_tensorrt.compile(traced_model, **compile_settings)
    print(f'Converted! used time {time.time() - s:.2f}s')

print("Saving ...")
torch.jit.save(model_trt, f'src/torch-tensortrt_fp16_{side}.trt')
time.sleep(20)

y = test(x, model, "model")
y_trt = test(x, model_trt, "model_trt")
print(f"Max error: {torch.max(torch.abs(y - y_trt))}")

结果:

torch.Size([1, 3, 512, 512])
Converting trt...
WARNING: [Torch-TensorRT] - Truncating weight (constant in the graph) from Float64 to Float32
...
WARNING: [Torch-TensorRT] - Truncating weight (constant in the graph) from Float64 to Float32
WARNING: [Torch-TensorRT TorchScript Conversion Context] - TensorRT was linked against cuBLAS/cuBLAS LT 11.6.3 but loaded cuBLAS/cuBLAS LT 11.5.1
WARNING: [Torch-TensorRT TorchScript Conversion Context] - TensorRT was linked against cuDNN 8.2.1 but loaded cuDNN 8.2.0
WARNING: [Torch-TensorRT TorchScript Conversion Context] - TensorRT was linked against cuBLAS/cuBLAS LT 11.6.3 but loaded cuBLAS/cuBLAS LT 11.5.1
WARNING: [Torch-TensorRT TorchScript Conversion Context] - TensorRT was linked against cuDNN 8.2.1 but loaded cuDNN 8.2.0
WARNING: [Torch-TensorRT] - TensorRT was linked against cuBLAS/cuBLAS LT 11.6.3 but loaded cuBLAS/cuBLAS LT 11.5.1
WARNING: [Torch-TensorRT] - TensorRT was linked against cuDNN 8.2.1 but loaded cuDNN 8.2.0
WARNING: [Torch-TensorRT] - TensorRT was linked against cuBLAS/cuBLAS LT 11.6.3 but loaded cuBLAS/cuBLAS LT 11.5.1
WARNING: [Torch-TensorRT] - TensorRT was linked against cuDNN 8.2.1 but loaded cuDNN 8.2.0
Converted! used time 198.42s
Saving ...
Testing model used 0.38s
torch.Size([1, 3, 1024, 1024])
Testing model_trt used 0.19s
torch.Size([1, 3, 1024, 1024])
Max error: 0.1142578125

批量测试优化效果:

import torch
import torch_tensorrt
from test.common import prepare, test, benchmark

side = 512
model, x = prepare(side)


torch.cuda.empty_cache()
model_trt = torch.jit.load(f'src/torch-tensortrt_fp16_{side}.trt')

y = test(x, model, "model")
y_trt = test(x, model_trt, "model_trt")
print("Max error", torch.max(torch.abs(y - y_trt)))


print('*'*100)
torch.cuda.empty_cache()
benchmark(model, x, warm_up=5, runs=50)

print('*'*100)
torch.cuda.empty_cache()
benchmark(model_trt, x, warm_up=5, runs=50)

结果:

torch.Size([1, 3, 512, 512])
WARNING: [Torch-TensorRT] - TensorRT was linked against cuBLAS/cuBLAS LT 11.6.3 but loaded cuBLAS/cuBLAS LT 11.5.1
WARNING: [Torch-TensorRT] - TensorRT was linked against cuDNN 8.2.1 but loaded cuDNN 8.2.0
WARNING: [Torch-TensorRT] - TensorRT was linked against cuBLAS/cuBLAS LT 11.6.3 but loaded cuBLAS/cuBLAS LT 11.5.1
WARNING: [Torch-TensorRT] - TensorRT was linked against cuDNN 8.2.1 but loaded cuDNN 8.2.0
Testing model used 0.46s
torch.Size([1, 3, 1024, 1024])
Testing model_trt used 0.25s
torch.Size([1, 3, 1024, 1024])
Max error tensor(0.1143, device='cuda:0', dtype=torch.float16)
****************************************************************************************************
Warm up ...
Start timing ...
Iteration 10/50, ave batch time 315.06 ms
Iteration 20/50, ave batch time 315.99 ms
Iteration 30/50, ave batch time 316.72 ms
Iteration 40/50, ave batch time 318.04 ms
Iteration 50/50, ave batch time 319.05 ms
Input shape: torch.Size([1, 12, 256, 256])
Output features size: torch.Size([1, 3, 1024, 1024])
Average batch time: 316.97 ms
****************************************************************************************************
Warm up ...
Start timing ...
Iteration 10/50, ave batch time 203.72 ms
Iteration 20/50, ave batch time 204.43 ms
Iteration 30/50, ave batch time 205.57 ms
Iteration 40/50, ave batch time 204.77 ms
Iteration 50/50, ave batch time 206.92 ms
Input shape: torch.Size([1, 12, 256, 256])
Output features size: torch.Size([1, 3, 1024, 1024])
Average batch time: 205.08 ms

加速效果和torch2trt类似。

Dynamic input shape
import time
import torch_tensorrt
import torch
from test.common import prepare, test

# Dynamic input shape
min_side = 128
opt_side = 256
max_side = 512

side = opt_side
model, x = prepare(side)

compile_settings = {
    "inputs": [
        torch_tensorrt.Input(
            min_shape=[1, 12, min_side//2, min_side//2],
            opt_shape=[1, 12, opt_side//2, opt_side//2],
            max_shape=[1, 12, max_side//2, max_side//2],
            dtype=torch.half)
    ],
    "truncate_long_and_double": True,
    "enabled_precisions": {torch.half}  # Run with FP16
}

with torch.no_grad():
    torch.cuda.empty_cache()
    traced_model = torch.jit.trace(model, x)
    torch.cuda.empty_cache()
    print('Converting trt...')
    s = time.time()
    model_trt = torch_tensorrt.compile(traced_model, **compile_settings)
    print(f'Converted! used time {time.time() - s:.2f}s')

print("Saving ...")
torch.jit.save(model_trt, f'src/torch-tensortrt_fp16_{min_side}-{max_side}.trt')
time.sleep(20)

y = test(x, model, "model")
y_trt = test(x, model_trt, "model_trt")
print(f"Max error: {torch.max(torch.abs(y - y_trt))}")

结果:

ssh://root@8.142.136.95:22/root/miniconda3/envs/ImageEnhance/bin/python -u /root/projects/imageenhance/test/torchrt.py
torch.Size([1, 3, 256, 256])
Converting trt...
WARNING: [Torch-TensorRT] - Truncating weight (constant in the graph) from Float64 to Float32
...
WARNING: [Torch-TensorRT] - Truncating weight (constant in the graph) from Float64 to Float32
WARNING: [Torch-TensorRT TorchScript Conversion Context] - TensorRT was linked against cuBLAS/cuBLAS LT 11.6.3 but loaded cuBLAS/cuBLAS LT 11.5.1
WARNING: [Torch-TensorRT TorchScript Conversion Context] - TensorRT was linked against cuDNN 8.2.1 but loaded cuDNN 8.2.0
WARNING: [Torch-TensorRT TorchScript Conversion Context] - TensorRT was linked against cuBLAS/cuBLAS LT 11.6.3 but loaded cuBLAS/cuBLAS LT 11.5.1
WARNING: [Torch-TensorRT TorchScript Conversion Context] - TensorRT was linked against cuDNN 8.2.1 but loaded cuDNN 8.2.0
WARNING: [Torch-TensorRT] - TensorRT was linked against cuBLAS/cuBLAS LT 11.6.3 but loaded cuBLAS/cuBLAS LT 11.5.1
WARNING: [Torch-TensorRT] - TensorRT was linked against cuDNN 8.2.1 but loaded cuDNN 8.2.0
WARNING: [Torch-TensorRT] - TensorRT was linked against cuBLAS/cuBLAS LT 11.6.3 but loaded cuBLAS/cuBLAS LT 11.5.1
WARNING: [Torch-TensorRT] - TensorRT was linked against cuDNN 8.2.1 but loaded cuDNN 8.2.0
Converted! used time 160.27s
Saving ...
Testing model used 0.13s
torch.Size([1, 3, 512, 512])
Testing model_trt used 0.08s
torch.Size([1, 3, 512, 512])
Max error: 0.05126953125

请添加图片描述
动态尺寸包含了固定尺寸的512宽度，但是转换时间更短160.27s<198.42s，模型文件更小136.1MB<153MB。
下面我们测试下时间

import torch
import torch_tensorrt
from test.common import prepare, test, benchmark

# min_side = int(5376 / 16 + 8)  # sc
# opt_side = int(6272 / 16 + 8)  # z1
# max_side = int(8000 / 16 + 8)  # xi_xun
# Dynamic input shape
min_side = 128
opt_side = 256
max_side = 512

side = max_side
model, x = prepare(side)


torch.cuda.empty_cache()
model_trt = torch.jit.load(f'../src/pretrained_models/real_esrgan/torch-tensortrt_fp16_{min_side}-{max_side}.trt')

y = test(x, model, "model")
y_trt = test(x, model_trt, "model_trt")
print("Max error", torch.max(torch.abs(y - y_trt)))


print('*'*100)
torch.cuda.empty_cache()
benchmark(model, x, warm_up=5, runs=50)

print('*'*100)
torch.cuda.empty_cache()
benchmark(model_trt, x, warm_up=5, runs=50)

结果：

torch.Size([1, 3, 512, 512])
WARNING: [Torch-TensorRT] - TensorRT was linked against cuBLAS/cuBLAS LT 11.6.3 but loaded cuBLAS/cuBLAS LT 11.5.1
WARNING: [Torch-TensorRT] - TensorRT was linked against cuDNN 8.2.1 but loaded cuDNN 8.2.0
WARNING: [Torch-TensorRT] - TensorRT was linked against cuBLAS/cuBLAS LT 11.6.3 but loaded cuBLAS/cuBLAS LT 11.5.1
WARNING: [Torch-TensorRT] - TensorRT was linked against cuDNN 8.2.1 but loaded cuDNN 8.2.0
Testing model used 0.43s
torch.Size([1, 3, 1024, 1024])
Testing model_trt used 0.21s
torch.Size([1, 3, 1024, 1024])
Max error tensor(0.1343, device='cuda:0', dtype=torch.float16)
****************************************************************************************************
Warm up ...
Start timing ...
Iteration 10/50, ave batch time 346.04 ms
Iteration 20/50, ave batch time 346.11 ms
Iteration 30/50, ave batch time 347.46 ms
Iteration 40/50, ave batch time 348.45 ms
Iteration 50/50, ave batch time 349.24 ms
Input shape: torch.Size([1, 12, 256, 256])
Output features size: torch.Size([1, 3, 1024, 1024])
Average batch time: 347.46 ms
****************************************************************************************************
Warm up ...
Start timing ...
Iteration 10/50, ave batch time 215.32 ms
Iteration 20/50, ave batch time 213.66 ms
Iteration 30/50, ave batch time 214.31 ms
Iteration 40/50, ave batch time 216.50 ms
Iteration 50/50, ave batch time 217.51 ms
Input shape: torch.Size([1, 12, 256, 256])
Output features size: torch.Size([1, 3, 1024, 1024])
Average batch time: 215.46 ms

加速效果和固定尺寸差不多: 215.46/347.46=0.62 205.08/316.97=0.65。

onnx-tensorrt

Installing onnx-tensorrt

安装步骤参考这个教程 ，已经很详细了，不在赘述。

Testing Conversion of RRDBNet

PyTorch -> ONNX

import torch
from test.common import prepare, test, show

side = 512
model, x = prepare(side)

torch.cuda.empty_cache()
with torch.no_grad():
    print('getting onnx...')
    y_onnx = torch.onnx._export(model, x, f'src/onnx_{side}.onnx', opset_version=11, export_params=True)
    show(y_onnx)

y = test(x, model, "model")
print("error", torch.max(torch.abs(y - y_onnx)))

ONNX -> TensorRT
使用命令行工具:

onnx2trt src/onnx_512.onnx -o src/onnx_512.trt_ -b 1 -d 16

d 16表示使用FP16精度

批量测试优化效果:

import time

import onnx
import torch
import onnx_tensorrt.backend as backend
import numpy as np
import tensorrt as trt
from test.common import prepare, test, show, benchmark

side = 512
model, x = prepare(side)

logger = trt.Logger()
with open(f'src/onnx_{side}.trt', "rb") as f, trt.Runtime(logger) as runtime:
    engine = runtime.deserialize_cuda_engine(f.read())
engine = backend.Engine(engine)


input_data = x[None].cpu().numpy()
print(input_data.shape)
s = time.time()
y_trt = engine.run(input_data)[0]
print('trt', time.time() - s)
y_trt = torch.from_numpy(np.array(y_trt)).cuda()
show(y_trt)

y = test(x, model, "model")
print("Max error", torch.max(torch.abs(y - y_trt)))

print('*'*100)
torch.cuda.empty_cache()
benchmark(model, x, warm_up=5, runs=50)

print('*'*100)
torch.cuda.empty_cache()
benchmark(engine.run, input_data, warm_up=5, runs=50)

结果:

torch.Size([1, 3, 512, 512])
WARNING: [Torch-TensorRT] - TensorRT was linked against cuBLAS/cuBLAS LT 11.6.3 but loaded cuBLAS/cuBLAS LT 11.5.1
WARNING: [Torch-TensorRT] - TensorRT was linked against cuDNN 8.2.1 but loaded cuDNN 8.2.0
WARNING: [Torch-TensorRT] - TensorRT was linked against cuBLAS/cuBLAS LT 11.6.3 but loaded cuBLAS/cuBLAS LT 11.5.1
WARNING: [Torch-TensorRT] - TensorRT was linked against cuDNN 8.2.1 but loaded cuDNN 8.2.0
Testing model used 0.40s
torch.Size([1, 3, 1024, 1024])
Testing model_trt used 0.19s
torch.Size([1, 3, 1024, 1024])
Max error tensor(0.1143, device='cuda:0', dtype=torch.float16)
****************************************************************************************************
Warm up ...
Start timing ...
Iteration 10/50, ave batch time 308.90 ms
Iteration 20/50, ave batch time 308.78 ms
Iteration 30/50, ave batch time 309.86 ms
Iteration 40/50, ave batch time 310.72 ms
Iteration 50/50, ave batch time 311.41 ms
Input shape: torch.Size([1, 12, 256, 256])
Output features size: torch.Size([1, 3, 1024, 1024])
Average batch time: 309.94 ms
****************************************************************************************************
Warm up ...
Start timing ...
Iteration 10/50, ave batch time 194.15 ms
Iteration 20/50, ave batch time 195.11 ms
Iteration 30/50, ave batch time 195.09 ms
Iteration 40/50, ave batch time 195.97 ms
Iteration 50/50, ave batch time 196.25 ms
Input shape: torch.Size([1, 12, 256, 256])
Output features size: torch.Size([1, 3, 1024, 1024])
Average batch time: 195.31 ms

加速效果和torch2trt以及Torch-TensorRT类似。

Conclusion

简单测试使用3种工具，发现对于PyTorch使用Torch-TensorRT 是最方便快捷的，支持多分辨率的输入输出。

snoopy_21
关注
  • 3
    点赞
  • 踩
  • 24
    收藏
    觉得还不错? 一键收藏
  • 2
    评论
RealESRGAN系列网络 C++&Win10&Tensorrt 部署教程
认真努力，做一只会飞的毛毛虫。
06-26 772
RealESRGAN系列网络 C++&Win10&Tensorrt 部署教程。
【2023 CCF 大数据与计算智能大赛】基于TPU平台实现超分辨率重建模型部署 基于Real-ESRGAN的TPU超分模型部署
算能开发者社区是一个为开发者提供信息传播、开发交流、技术分享的交流空间。
01-05 1355
也就是说，这个下采样的模糊核是预先定义好的。模型部署使用的是算能提供的C++版本的 SophonSDK中的bmcv等加速库进行图像加载、数据前处理、模型推理、数据后处理等操作。SophonSDK是算能科技基于其自主研发的 AI芯片所定制的深度学习 SDK，涵盖了神经网络推理阶段所需的模型优化、高效运行支持等能力，为深度学习应用开发和部署提供易用、高效的全栈式解决方案。模型部署使用的是算能提供的C++版本的 SophonSDK中的bmcv等加速库进行图像加载、数据前处理、模型推理、数据后处理等操作[3]。
一行代码加速Pytorch推理速度6倍
专注于人工智能领域的小何尚
04-12 933
PyTorch 有一个名为 PyTorch Hub 的模型存储库，它是常见模型的高质量实现的来源。我们可以从那里获得在 ImageNet 上预训练的 ResNet-50 模型。在本教程中，我们介绍了使用 Torch-TensorRT 为 ResNet-50 模型编译 TorchScript 模型的完整过程，并测试了优化对性能的影响。使用 Torch-TensorRT，我们在 NVIDIA V100 GPU 上观察到 FP32 的加速为1.42x，FP16 的加速为5.4x。
TensorRT-5.1.5.0-YOLOv3
baidu_40840693的博客
07-14 1081
接上一个博客 TensorRT-5.1.5.0-SSD https://blog.csdn.net/baidu_40840693/article/details/95642055 首先： upsample层 密码bwrd https://pan.baidu.com/s/13GpoYoqKSCeFX0m0ves_fQ 然后： https://github.com/lewes636...
tensorrt动态输入分辨率尺寸
jacke121的专栏
03-27 3461
本文只有 tensorrt python部分涉动态分辨率设置，没有c++的。 目录 pytorch转onnx： onnx转tensorrt： python tensorrt推理： 知乎博客也可以参考： tensorrt动态输入（Dynamic shapes） - 知乎 记录此贴的原因有两个：1.肯定也有很多人需要 。2.就我搜索的帖子没一个讲的明明白白的，官方文档也不利索，需要连蒙带猜。话不多少，直接上代码。 以pytorch转onnx转tensorrt为例，动态shape是图像的长..
深度学习模型试跑(十五):Real-ESRGAN(VS2019 trt推理部署)
qq_33642342的博客
09-08 3702
超分辨率修复、重建是采用低分辨率（LR）输入并将其提高到高分辨率的任务，具体原理可以参考paddlegan原理介绍。对于在linux上如何实现trt加速该网络，这里有一篇文章详细记录了过程。
转动态分辨率tensorRT的pytorch代码修改
WX_Chen的博客
09-01 443
1、条件分支删除 2、view不能输入-1 3、torch.tensor加@torch.jit.script修饰，入参必须是tensor类型，不能用.to(device) 4、不能tensor转python类型，改用//运算 5、reshape改为view，要加.contiguous() 6、不能用max()或min() 遇到问题： Get cuda error during getBestTactic: PWN(Sigmoid_3299 解决： try a small batch size and
realesrgan-ncnn-vulkan批量处理图片的Powershell脚本分享
Sakurananasang的博客
03-08 1373
realesrgan-ncnn-vulkan批量处理图片的Powershell脚本分享
Real-ESRGAN 深度学习超分辨率重建工具
02-09
总结来说，Real-ESRGAN是一款基于深度学习的高效超分辨率重建工具，通过其预训练模型，用户能够轻松地提升图像的分辨率，达到更高质量的视觉效果。它利用了深度学习技术，特别是生成对抗网络，使得即使在没有GPU加速...
基于Real-ESRGAN的图片AI超分辨率安卓APP完整源码+说明.zip
03-26
1、该资源内项目代码都是经过测试运行成功，功能正常的情况下才上传的，请放心下载使用。 2、适用人群：主要针对计算机相关专业(如计科、信息安全、数据科学与大数据技术、人工智能、通信、物联网、数学、电子信息等...
Real-ESRGAN-master
最新发布
05-08
《Real-ESRGAN-master: 探索深度学习在图像超分辨率中的应用》 在现代数字图像处理领域，超分辨率重建（Super-Resolution）技术已经成为一个重要的研究热点。Real-ESRGAN-master项目，正如其名，是针对真实世界图像...
基于Real-ESRGAN的图片AI超分辨率安卓APP.zip
02-06
我们深入探讨了深度学习的基本原理、神经网络的应用、自然语言处理、语言模型、文本分类、信息检索等领域。更有深度学习、机器学习、自然语言处理和计算机视觉的实战项目源码，助您从理论走向实践，如果您已有一定...
基于Real-ESRGAN的图片AI超分辨率安卓APP源码+项目说明.zip
03-16
【资源说明】 1、该资源包括项目的全部源码，下载可以直接使用！ 2、本项目适合作为计算机、数学、电子信息等专业的课程设计、期末大作业和毕设项目，作为...基于Real-ESRGAN的图片AI超分辨率安卓APP源码+项目说明.zip
【Python&图像超分】Real-ESRGAN图像超分模型（超分辨率重建）详细安装和使用教程
Remote Sensing迷途小书童
11-14 9641
图像超分是一种图像处理技术，旨在提高图像的分辨率，使其具有更高的清晰度和细节。这一技术通常用于图像重建、图像恢复、图像增强等领域，可以帮助我们更好地理解和利用图像信息。今天给大家介绍一下腾讯ARC实验室发布的一个图像超分辨率模型Real-ESRGAN，同时奉上详细的安装使用教程。
自己写的超分软件 realesrgan超分工具脚本版
2301_79359407的博客
11-08 156
在realesrgan超分工具脚本版当中，你可以选择为照片，动漫图片，动漫视频超分，你也可以在此工具里面更改视频分辨率。这些操作都是极其简单的，不用掌握real - esrgan 的命令行用法，轻松使用real - esrgan 的模型为图片以及视频进行超分。另外，如果有动漫视频超分的需求，请尽量把软件安装在大容量硬盘里，因为视频的解构与重构可能会产生大量缓存文件。软件支持大量图片批量超分，视频批量超分，速度取决于您的硬件。自己写的超分软件 realesrgan超分工具。视频超分，右原视频，左超分。
Torch-TensorRT安装、简单使用；查看显卡实时情况nvitop、nvidia-smi、gpustat
weixin_42357472的博客
06-17 4103
***因为需要各种tensort、cuda等环境配置，所以可以直接安装索引环境都已经搭好的docker 参考：https://catalog.ngc.nvidia.com/orgs/nvidia/containers/pytorch ****** docker images：nvcr.io/nvidia/pytorch 22.05-py3 14.6GB*****运行： 进入docker： 2、简单使用Torch-TensorRT 参考：https://
【pytorch，onnx，bug解决】pytorch的op：pixel unshuffle转到onnx的op：SpaceToDepth
一个CV小白的博客
06-28 1342
pytorch的`F.pixel_unshuffle` 转onnx 的 `SpaceToDepth`，
pytorch的 F.pixel_unshuffle函数理解，利用卷积来实现pixel_unshuffle, yolov5中focus操作和pixel_unshuffle是否一样？
tony365的博客
04-03 453
b3是yolov5中用到的focus操作，上述代码我用卷积实现了pixel_unshuffle操作， focus操作同理当然也可以利用卷积实现。pixel_unshuffle示意图：和focus不同的是最后的输出的通道排列顺序有差异。结果显示：b = b1 =b2 但是和b3不太一样.
TensorRT部署YOLOv5(03)-TensorRT介绍
abcxy2161的博客
01-09 832
TensorRT是本专栏中最重要的内容，绝大多数内容将围绕TensorRT来展开，本文对TensorRT进行一个基本的介绍，让不熟悉TensorRT的读者能够对TensorRT是什么，如何使用它有一个较为全面的认识Nvidia TensorRT是一个用于Nvidia GPU上高性能机器学习推理的SDK，对开发者屏蔽了模型在GPU上推理运行的CUDA计算细节，用户只需要通过一套简介易用的Python/C++ API接口，即可方便的将模型在GPU上进行加速推理。
C++ tensorRT部署real-esrgan onnx
05-26
要在C++中使用TensorRT部署Real-ESRGAN ONNX模型，可以按照以下步骤进行： 1. 安装TensorRT并设置环境变量，确保能够在C++代码中引用TensorRT头文件和库文件。 2. 将Real-ESRGAN ONNX模型转换为TensorRT格式。可以使用TensorRT提供的ONNX Parser API将ONNX模型转换为TensorRT格式。转换后的模型可以使用TensorRT的C++ API加载和推理。 3. 在C++代码中加载转换后的TensorRT模型，并创建推理引擎和执行上下文。推理引擎可以使用TensorRT的Builder API创建，执行上下文可以使用TensorRT的Execution API创建。 4. 准备输入数据并将其传递给推理引擎的执行上下文。可以使用TensorRT的HostDeviceMem类来管理输入和输出数据的内存，以确保正确的数据传输和内存访问。 5. 执行推理并获取输出结果。可以使用TensorRT的ExecutionContext类来执行推理并获取输出结果。 下面是一个简单的C++代码示例，用于加载和推理Real-ESRGAN ONNX模型： ```c++ #include <NvInfer.h> #include <NvOnnxParser.h> #include <iostream> using namespace nvinfer1; using namespace nvonnxparser; int main() { // 1. Create a TensorRT builder and network IBuilder* builder = createInferBuilder(gLogger); INetworkDefinition* network = builder->createNetworkV2(0U); // 2. Parse the ONNX model and create a TensorRT engine IParser* parser = createParser(*network, gLogger); parser->parseFromFile("real_esrgan.onnx", 0); builder->setMaxBatchSize(1); builder->setMaxWorkspaceSize(1 << 30); ICudaEngine* engine = builder->buildCudaEngine(*network); // 3. Create an execution context and allocate memory for input and output buffers IExecutionContext* context = engine->createExecutionContext(); void* inputBuffers[1]; void* outputBuffers[1]; int inputIndex = engine->getBindingIndex("input_1"); int outputIndex = engine->getBindingIndex("Identity"); Dims inputDims = engine->getBindingDimensions(inputIndex); Dims outputDims = engine->getBindingDimensions(outputIndex); int inputSize = sizeof(float) * inputDims.d[0] * inputDims.d[1] * inputDims.d[2] * inputDims.d[3]; int outputSize = sizeof(float) * outputDims.d[0] * outputDims.d[1] * outputDims.d[2] * outputDims.d[3]; cudaMalloc(&inputBuffers[0], inputSize); cudaMalloc(&outputBuffers[0], outputSize); // 4. Prepare input data and copy to input buffer float* inputData = new float[inputDims.d[0] * inputDims.d[1] * inputDims.d[2] * inputDims.d[3]]; // fill input data with your own data here cudaMemcpy(inputBuffers[0], inputData, inputSize, cudaMemcpyHostToDevice); // 5. Run inference and get output data context->executeV2(inputBuffers); float* outputData = new float[outputDims.d[0] * outputDims.d[1] * outputDims.d[2] * outputDims.d[3]]; cudaMemcpy(outputData, outputBuffers[0], outputSize, cudaMemcpyDeviceToHost); // do something with the output data here // 6. Clean up delete[] inputData; delete[] outputData; cudaFree(inputBuffers[0]); cudaFree(outputBuffers[0]); context->destroy(); engine->destroy(); network->destroy(); builder->destroy(); return 0; } ``` 注意，此示例仅用于演示目的，并不包含完整的错误处理和内存管理代码。在实际应用中，应该根据实际情况进行适当的修改和完善。

“相关推荐”对你有帮助么？

  • 非常没帮助
  • 没帮助
  • 一般
  • 有帮助
  • 非常有帮助
提交
评论 2
添加红包

请填写红包祝福语或标题

个

红包个数最小为10个

元

红包金额最低5元

当前余额3.43元 前往充值 >
需支付：10.00元
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付元
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明：

1.余额是钱包充值的虚拟货币，按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载，可以购买VIP、付费专栏及课程。

余额充值