直接将以下脚本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

Logo

GitCode 天启AI是一款由 GitCode 团队打造的智能助手,基于先进的LLM(大语言模型)与多智能体 Agent 技术构建,致力于为用户提供高效、智能、多模态的创作与开发支持。它不仅支持自然语言对话,还具备处理文件、生成 PPT、撰写分析报告、开发 Web 应用等多项能力,真正做到“一句话,让 Al帮你完成复杂任务”。

更多推荐