还在为模型加速推理发愁吗？不如看看这篇吧。手把手教你把pytorch模型转化为TensorRT，加速推理

到这里第一种写法就完成了，是不是很简单，接下来看第二种写法。

第二种推理写法

---

新建onnx.py脚本，加入以下代码：

import onnxruntime

class ONNXModel():

def init(self, onnx_path):

“”"

:param onnx_path:

“”"

self.onnx_session = onnxruntime.InferenceSession(onnx_path,providers=[‘TensorrtExecutionProvider’, ‘CUDAExecutionProvider’, ‘CPUExecutionProvider’])

self.input_name = self.get_input_name(self.onnx_session)

self.output_name = self.get_output_name(self.onnx_session)

print(“input_name:{}”.format(self.input_name))

print(“output_name:{}”.format(self.output_name))

def get_output_name(self, onnx_session):

“”"

output_name = onnx_session.get_outputs()[0].name

:param onnx_session:

:return:

“”"

output_name = []

for node in onnx_session.get_outputs():

output_name.append(node.name)

return output_name

def get_input_name(self, onnx_session):

“”"

input_name = onnx_session.get_inputs()[0].name

:param onnx_session:

:return:

“”"

input_name = []

for node in onnx_session.get_inputs():

input_name.append(node.name)

return input_name

def get_input_feed(self, input_name, image_numpy):

“”"

input_feed={self.input_name: image_numpy}

:param input_name:

:param image_numpy:

:return:

“”"

input_feed = {}

for name in input_name:

input_feed[name] = image_numpy

return input_feed

def forward(self, image_numpy):

输入数据的类型必须与模型一致,以下三种写法都是可以的

scores, boxes = self.onnx_session.run(None, {self.input_name: image_numpy})

scores, boxes = self.onnx_session.run(self.output_name, input_feed={self.input_name: iimage_numpy})

input_feed = self.get_input_feed(self.input_name, image_numpy)

scores = self.onnx_session.run(self.output_name, input_feed=input_feed)

return scores

调用onnx.py实现推理，新建test_onnx1.py插入代码：

import os, sys

sys.path.append(os.getcwd())

import numpy as np

import torchvision.transforms as transforms

from PIL import Image

from models.onnx import ONNXMode