Pytorch 模型tensorrt部署v1

0. 目录

更新 2020.05

onnx将模型转为trt或者其他inference是一种比较通用的方式，博主测试过程中发现，onnx转换出来的模型会稍慢与纯trt api 写出来的模型，nvidia官方也有说到过onnx存在效率非最优以及微小数值差异的问题。
为此，强烈建议大家了解下 torch2trt这个框架，可以无痛转换模型并保持最好的效率。github.torch2trt
针对常见的分类，检测，分割问题是非常好用的，同时支持定义converter，拓展能力也还不错。强推荐。日后博主会找几个问题进行更具体的介绍以及尝试。

更新2020.06

1. 对 resnest 网络进行加速

项目全部代码:https://github.com/musicbeer/resnest-tensorRT
resnest网络刷爆多个CV任务，虽然在review时候有些争议，但实际项目中好用就是硬道理。这里可以尝试下利用torch2trt 对该网络结构进行加速，同时测试一下resnet速度到底怎么样。
首先我们需要装一下torch2trt这个包：

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

然后加载模型文件并做前向测试

from test.resnest import resnest50
import torch
from torch2trt import torch2trt 
from timeit import default_timer as timer
from torchvision import models
net = resnest50(pretrained=True).cuda().eval()
x = torch.ones((1, 3, 224, 224))
with torch.no_grad():
    y_pt = net(x.cuda())[0]

在这里会报warning，说tensor的transpose 操作不支持，但是仔细看了converters中明明有这个操作。最后的解决方法只需要修改splat.py 中的一行代码即可

class rSoftMax(nn.Module):
    def __init__(self, radix, cardinality):
        super().__init__()
        self.radix = radix
        self.cardinality = cardinality

    def forward(self, x):
        batch = x.size(0)
        if self.radix > 1:
            x = torch.transpose(x.view(batch, self.cardinality, self.radix, -1),1,2)   #.permute(0,2, 1, 3) ***
            x = F.softmax(x, dim=1)
            x = x.reshape(batch, -1)
        else:
            x = torch.sigmoid(x)
        return x

不出意外的话就能够完成转换了。为了做速度对比，博主对resnet50以及resnet101 做了同样的测试，测试结果如下：

python test.py 

----------------------resnest 50 test---------------------
show first 10 output:
pytorch: tensor([-0.8604,  0.3855,  0.3289,  0.1835,  0.7670, -0.2655, -0.4620, -0.5909,
        -0.9840,  0.3266], device='cuda:0')
tensorRT: tensor([-0.8604,  0.3855,  0.3289,  0.1835,  0.7670, -0.2655, -0.4620, -0.5909,
        -0.9840,  0.3266], device='cuda:0')


error: tensor(0.0024, device='cuda:0')


3x224x224 image inference time:
pytorch: 0.016652769159991296
tensorRT: 0.009595841290429235
----------------------resnet 50 test---------------------
show first 10 output:
pytorch: tensor([-0.3081,  0.0798, -1.1900, -1.4837, -0.5136,  0.3683, -2.1639, -0.8705,
        -1.8812, -0.1608], device='cuda:0')
tensorRT: tensor([-0.3081,  0.0798, -1.1900, -1.4837, -0.5136,  0.3683, -2.1639, -0.8705,
        -1.8812, -0.1608], device='cuda:0')


error: tensor(0.0011, device='cuda:0')


3x224x224 image inference time:
pytorch: 0.01067693106131628
tensorRT: 0.005997186410240829
----------------------resnet 101 test---------------------
show first 10 output:
pytorch: tensor([-0.9190, -0.1303, -1.0331, -1.5215, -0.5173,  0.1020, -2.0780, -0.8331,
        -1.4745, -0.4834], device='cuda:0')
tensorRT: tensor([-0.9190, -0.1303, -1.0331, -1.5215, -0.5173,  0.1020, -2.0780, -0.8331,
        -1.4745, -0.4834], device='cuda:0')


error: tensor(0.0013, device='cuda:0')


3x224x224 image inference time:
pytorch: 0.019209207659587266
tensorRT: 0.011350022009573878




resnest50 trt speedup 16ms->9.6ms, resnet50 11ms->6.1ms,resnet101 19ms->11ms,so resnest50 is faster than resnet101.

按照论文中说的，resnest56在多个项目的效果超过resnet101，那么可以说作者确实找到了一种更优秀的网络结构。

博主推荐以下的内容不用看了，如果你非要看，那你看。

2. 主要内容

项目全部代码:https://github.com/musicbeer/pytorch-tensorrt

为了提高神经网络模型的推理速度,Nvidia 推出了 Tensorrt 网络前向加速器. 简单的来说,训练好一个神经网络模型,只需要将模型通过tensorrt部署就会获得更快的推理速度. Tensorrt中支持FP32,FP16,INT8数据类型,其中FP32与FP16在精度上的损失很小,INT8量化后的模型精度上的损失稍大,但一般的分类模型可以使用.当然我们可以根据不同需求以及模型效果进行选择.

本文中使用的方法是 首先将pytorch模型转换到onnx model, 然后通过tensorrt中提供的onnxparse将onnx 模型转到 trt 模型.

2.1 准备阶段

深度学习环境配置相对繁琐,强烈推荐docker. Docker hub 有现成的tensorrt基础镜像,直接pull下来用就好.

了解docker, tensorrt, ONNX.

感谢https://github.com/modricwang/Pytorch-Model-to-TensorRT,因为绝大部分代码都是在这个工程上面修改的.

2.2 主要代码

1. 首先import 一些需要的库

import torch
from torch.autograd import Variable
import tensorrt as trt
from tensorrt.parsers import onnxparser
import pycuda.driver as cuda
import pycuda.gpuarray as gpuarray
import pycuda.autoinit
from trt_engine import trt_engine
import os
import numpy as np
from argparse import ArgumentParser
from resnet import resnet50
from image_reader import read_image_chw
import calib as calibrator
import time

2. 定义各个步骤的模型名称:

args = ArgumentParser().parse_args()
args.input_size = 224
args.input_channel = 3
args.fc_num = 1000
args.batch_size = 4
args.onnx_model_name = "resnet50.onnx"
args.trt32_model_name = "resnet50_32.trt"

3. 接下来,通过torch.onnx.export 函数将加载好的模型转换到ONNX模型 resnet50.onnx,是不是简单粗暴,代码中的dummy_input中明确了网络输入数据的具体大小.顺利的话会直接生成resnet50.onnx文件,值得注意的是onnx 与tensorrt 支持的layer有限,如果出现不能识别的layer或者参数的话转换任务就不能完成.在resnet.py中,x = x.reshape(x.size(0), -1)需修改为x = x.view(int(x.shape[0]), -1)就不会报错了.

model=resnet50(pretrained=True)
model.cuda()

# Translate Pytorch Model into Onnx Model
dummy_input = Variable(torch.randn(args.batch_size, args.input_channel, \
                                   args.input_size, args.input_size, device='cuda'))
output_names = ["output"]
torch.onnx.export(model, dummy_input, args.onnx_model_name, verbose=False,
                  output_names=output_names)

4. 最关键的一步,我们需要将onnx模型转换到tensorrt.一般的首次成功加载了trt engine之后会将序列化后的文件进行存储.这样在部署的时候就可以直接读存储后的文件来加快模型加载的速度.trt支持动态batch_size.在构建的时候指定最大的batchsize即可.

def onnx_2_float32():
    apex = onnxparser.create_onnxconfig()
    apex.set_model_file_name(args.onnx_model_name)
    apex.set_model_dtype(trt.infer.DataType.FLOAT)
    apex.set_print_layer_info(False)
    trt_parser = onnxparser.create_onnxparser(apex)

    data_type = apex.get_model_dtype()
    onnx_filename = apex.get_model_file_name()
    trt_parser.parse(onnx_filename, data_type)
    trt_parser.convert_to_trtnetwork()
    trt_network = trt_parser.get_trtnetwork()

    G_LOGGER = trt.infer.ConsoleLogger(trt.infer.LogSeverity.ERROR)
    builder = trt.infer.create_infer_builder(G_LOGGER)
    builder.set_max_batch_size(16)
    engine = builder.build_cuda_engine(trt_network)
    modelstream = engine.serialize()
    trt.utils.write_engine_to_file(args.trt32_model_name, modelstream)
    engine.destroy()
    builder.destroy()

ok 执行函数onnx_2_float32(),resnet50_32.trt文件生成就说明转换工作已经完成了.

5. 测试

   import torch
   from torch.autograd import Variable
   import tensorrt as trt
   from tensorrt.parsers import onnxparser
   import pycuda.driver as cuda
   import pycuda.gpuarray as gpuarray
   import pycuda.autoinit
   
   import os
   import numpy as np
   from argparse import ArgumentParser
   
   from resnet import resnet50
   import time
   import cv2
   from torchvision import transforms
   from trt_engine import trt_engine
   import pdb
   MEAN = [0.485, 0.456, 0.406]
   STD = [0.229, 0.224, 0.225]
   
   transform = transforms.Compose([
       transforms.ToTensor(),
       transforms.Normalize(mean=MEAN, std=STD),
   ])
   def read_image_chw(img_path, width, height):### tensorrt 数据加载
       img = cv2.imread(img_path)
       img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
   
       img = cv2.resize(img, (width, height))
   
       im = transform(img)
       im = im.numpy()
       return im
   args = ArgumentParser().parse_args()
   args.input_size = 224
   args.input_channel = 3
   args.fc_num = 1000
   args.batch_size = 4
   args.trt_model_name = "resnet50_32.trt"
   args.testdir='data'
   engine = trt_engine('resnet',args.trt_model_name).build_engine()###将上一步生成的resnet50_32.trt 加载到trt engine中
   # 这里一般将testdir设置为训练pytorch模型时候的验证或者测试集合即可,本测试代码可以统计准确率,本代码只测试了一张图像.并输出了最后1000维特征的前5个进行输出.
   classes=os.listdir(args.testdir)
   total = 0
   correct = 0
   # speed test 
   
   model = resnet50(pretrained=True)
   model.cuda()
   model.eval()
   trans=transforms.Compose([transforms.Resize([224,224]),transforms.ToTensor(),
                                                transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),])
   classes=os.listdir(args.testdir)
   total = 0
   correct = 0
   start=time.time()
   for i in range(len(classes)):
       images=os.listdir(os.path.join(args.testdir,classes[i]))
       for image in images:
           gt=int(i)
           img = cv2.imread(os.path.join(args.testdir,classes[i],image))
           img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
           img = cv2.resize(img, (args.input_size, args.input_size))
           img = transform(img)
           img=torch.unsqueeze(img, 0)## 保持输入与tensorrt输入一致,这里也可以测试cv2与PIL 加载图像 resize的效率.
   
           batch_x=Variable(img)
           batch_x=batch_x.cuda()
           out = model(batch_x)
           print(out[0][:5])
           _,pred=torch.max(out.data, 1)
           pred = int(pred)
           # print('gt:',classes[i],'pred:',classes[pred])
           # if pred == gt:
           #     correct += 1
           # total += 1
   end=time.time()
   print('pytorch:',end-start)
   start=time.time()
   for i in range(len(classes)):
       images=os.listdir(os.path.join(args.testdir,classes[i]))
       for image in images:
           gt=int(i)
           img = read_image_chw(os.path.join(args.testdir,classes[i],image),
                   args.input_size, args.input_size)
           output = engine.infer(img)
           print(np.squeeze(output[0])[:5])
           conf, pred = torch.Tensor(np.squeeze(output[0])).topk(1, dim=0)
           pred = int(pred.data[0])
           # print('gt:',classes[i],'pred:',classes[pred])
           # if pred == gt:
           #     correct += 1
           # total += 1
   end=time.time()
   print('trt:',end-start)
   
   

   对比一下最终结果.是相同的.而且速度有一定的提升.