5、自动获取onnx模型的输入输出节点，并进行前向推理

最新推荐文章于 2025-04-02 19:27:50 发布

原创最新推荐文章于 2025-04-02 19:27:50 发布 · 1.7k 阅读

5 ·

CC 4.0 BY-SA版权

文章标签：

#深度学习 #tensorflow #人工智能

python笔记整理专栏收录该内容

9 篇文章

订阅专栏

该博客介绍了如何利用ONNXRuntime加载和执行一个ONNX模型。首先创建了一个ONNXModel类，初始化时会读取ONNX模型并获取输入和输出名称。`get_input_feed`方法用于构建输入数据的字典，`forward`方法则执行模型推理。在示例中，随机生成一个4维张量并进行推理，输出了张量的形状。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

import onnx
import onnxruntime
class 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_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_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_tensor
        return input_feed

    def 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) #根据数据决定输出参数的个数，此处为两个（scores,boxes）
        return scores,boxes
if __name__=='__main__':
    abc = ONNXModel('your onnx_model path ')
    abc ,box= abc.forward(np.random.randn(1, 3, 640, 640).astype(np.float32))
    print(abc[1].shape)