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

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