onnx模型推理(python)

onnx模型推理(python)

以下ONNX一个检测模型的推理过程,其他模型稍微修改即可

# -*-coding: utf-8 -*-import os, syssys.path.append(os.getcwd()) import onnxruntime import onnxclass ONNXModel():def __init__(self, onnx_path):""":param onnx_path:"""self.onnx_session = onnxruntime.InferenceSession(onnx_path)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_namedef 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_namedef get_input_feed(self, input_name, image_tensor):"""input_feed={self.input_name: image_tensor}:param input_name::param image_tensor::return:"""input_feed = {}for name in input_name:input_feed[name] = image_tensorreturn input_feeddef forward(self, image_tensor):'''image_tensor = image.transpose(2, 0, 1)image_tensor = image_tensor[np.newaxis, :]onnx_session.run([output_name], {input_name: x}):param image_tensor::return:'''# 输入数据的类型必须与模型一致,以下三种写法都是可以的# scores, boxes = self.onnx_session.run(None, {self.input_name: image_tensor})# scores, boxes = self.onnx_session.run(self.output_name, input_feed={self.input_name: image_tensor})input_feed = self.get_input_feed(self.input_name, image_tensor)scores, boxes = self.onnx_session.run(self.output_name, input_feed=input_feed)return scores, boxes

 

总结

以上是生活随笔为你收集整理的onnx模型推理(python)的全部内容，希望文章能够帮你解决所遇到的问题。

如果觉得生活随笔网站内容还不错，欢迎将生活随笔推荐给好友。