获取 onnx 模型的输入输出信息 Python 脚本

氵文大师

已于 2022-08-04 23:37:51 修改

阅读量7.6k

点赞数 4

分类专栏： onnx 文章标签： python 开发语言深度学习

于 2022-08-04 23:17:00 首次发布

本文链接：https://blog.csdn.net/HaoZiHuang/article/details/126168132

版权

onnx 专栏收录该内容

5 篇文章

订阅专栏

这段脚本展示了如何使用onnxruntime库获取ONNX模型的输入和输出信息。当模型输入shape固定时，例如[1,3,224,224]，输出为固定的形状；若输入shape可变，如[None,3,None,None]，则表示模型接受不同尺寸的输入。

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

直接将以下脚本onnx模型路径 onnx_path 修改即可：

from pprint import pprint
import onnxruntime

onnx_path = "originalpool/output.onnx"
# onnx_path = "custompool/output.onnx"

provider = "CPUExecutionProvider"
onnx_session = onnxruntime.InferenceSession(onnx_path, providers=[provider])

print("----------------- 输入部分 -----------------")
input_tensors = onnx_session.get_inputs()  # 该 API 会返回列表
for input_tensor in input_tensors:         # 因为可能有多个输入，所以为列表
    
    input_info = {
        "name" : input_tensor.name,
        "type" : input_tensor.type,
        "shape": input_tensor.shape,
    }
    pprint(input_info)

print("----------------- 输出部分 -----------------")
output_tensors = onnx_session.get_outputs()  # 该 API 会返回列表
for output_tensor in output_tensors:         # 因为可能有多个输出，所以为列表
    
    output_info = {
        "name" : output_tensor.name,
        "type" : output_tensor.type,
        "shape": output_tensor.shape,
    }
    pprint(output_info)

值得说明的是，如果onnx模型的输入shape是固定的，该脚本的输出是：

'----------------- 输入部分 -----------------'
{'name': 'x', 
 'shape': [1, 3, 224, 224], 
 'type': 'tensor(float)'}
'----------------- 输出部分 -----------------'
{'name': 'bilinear_interp_v2_7.tmp_0',
 'shape': [1, 2, 224, 224],
 'type': 'tensor(float)'}

值得说明的是，如果onnx模型的输入shape是非固定的，该脚本的输出是：

----------------- 输入部分 -----------------
{'name': 'x', 
 'shape': [None, 3, None, None], 
 'type': 'tensor(float)'}
----------------- 输出部分 -----------------
{'name': 'bilinear_interp_v2_7.tmp_0',
 'shape': [None, 2, None, None],
 'type': 'tensor(float)'}

shape 中有 None