tersorrt安装_pytorch/mxnet模型tensorrt部署

本文用于记录pytorch/mxnet模型使用tersorrt的整个流程以及遇到的坑。

tensorrt支持TensorFlow的uff和onnx以及自定义模型的推理加速，对于pytorch有第三方接口torch2trt项目，但是这个需要定义好模型在加入，不能把模型和tensorrt分离

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

部署的时候还依赖pytorch环境，就没尝试。

mxnet官方是有接口直接转tensorrt的，

arg_params.update(aux_params)

all_params = dict([(k, v.as_in_context(mx.gpu(0))) for k, v in arg_params.items()])

executor = mx.contrib.tensorrt.tensorrt_bind(sym, ctx=mx.gpu(0), all_params=all_params,data=batch_shape, grad_req='null', force_rebind=True)

y_gen = executor.forward(is_train=False, data=input)

y_gen[0].wait_to_read()

这个也没有尝试，主要还是想部署时分离，只用tensorrt环境，不需要装深度学习全家桶

pytorch和mxnet转换为onnx的模型官方都有接口和文档，使用方法也很简单

#mxnet转onnx

sym = './resnet-50-symbol.json'

params = './resnet-50-0000.params'

input_shape = (1, 3, 224, 224)

onnx_file = './resnet-50.onnx'

converted_model_path = onnx_mxnet.export_model(sym, params, [input_shape], np.float32, onnx_file)

#pytorch转onnx

import torch

import torchvision

dummy_input = torch.randn(10, 3, 224, 224, device='cuda')

model = torchvision.models.alexnet(pretrained=True).cuda()

# Providing input and output names sets the display names for values

# within the model's graph. Setting these does not change the semantics

# of the graph; it is only for readability.

#

# The inputs to the network consist of the flat list of inputs (i.e.

# the values you would pass to the forward() method) followed by the

# flat list of parameters. You can partially specify names, i.e. provide

# a list here shorter than the number of inputs to the model, and we will

# only set that subset of names, starting from the beginning.

input_names = [ "actual_input_1" ] + [ "learned_%d" % i for i in range(16) ]

output_names = [ "output1" ]

torch.onnx.export(model, dummy_input, "a