【工具】onnx模型结构信息查看方式：netron、onnxruntime和onnx

前言

将模型转为onnx格式进行部署时，对模型结构不是特别清楚，尤其是模型的输入输出结构时，通常使用一些工具进行查看。

netron

netron是一个非常直观的模型可视化工具，支持多种模型格式，包括ONNX。Netron允许用户以图形化方式查看模型的层级结构、输入输出节点及其详细信息。用户可以在浏览器中直接使用Netron，或者下载桌面版应用程序进行本地查看。

地址：https://netron.app/

上图可以看到，模型的输入为[1，3，640，640],输出为[1，84，8400]

• [1，3，640，640]:batch为1，输入图片为640 * 640
• [1，84，8400]:batch为1,输出向量为84（代表x, y, w, h,cls类别数80），检测框数目为8400。

onnxruntime

onnxruntime作为ONNX的官方运行时库，ONNX Runtime不仅支持模型推理，还提供了对模型结构进行基本检查的功能。用户可以通过onnxruntime的API加载模型并查询输入输出信息。
代码：

import onnxruntime as ort

providers = ["CPUExecutionProvider"]
session_options = ort.SessionOptions()
session_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL

session = ort.InferenceSession(
    'yolov8.onnx',
    session_options=session_options,
    providers=providers
)

inputs_info = session.get_inputs()
print("模型的输入信息:")
for input in inputs_info:
    print(f"Name: {input.name}, Shape: {input.shape}, Type: {input.type}")

outputs_info = session.get_outputs()
print("\n模型的输出信息:")
for output in outputs_info:
    print(f"Name: {output.name}, Shape: {output.shape}, Type: {output.type}")

输出：

模型的输入信息:
Name: images, Shape: [1, 3, 640, 640], Type: tensor(float)

模型的输出信息:
Name: output0, Shape: [1, 10, 8400], Type: tensor(float)

ONNX

ONNX提供了Python API，用于加载和操作ONNX模型。通过ONNX Python API，用户可以编写脚本以程序化方式检查模型的详细结构，包括输入输出节点及其数据类型。

代码：

import onnx
model = onnx.load('yolov8.onnx')
for imp in model.opset_import:
   print(f"打印OpSet导入信息: {imp.domain}, Version: {imp.version}")
# 验证模型是否有效
onnx.checker.check_model(model)
print("模型的输入信息:")
for input in model.graph.input:
    print(f"Name: {input.name}, Type: {input.type}")
print("\n模型的输出信息:")
for output in model.graph.output:
    print(f"Name: {output.name}, Type: {output.type}")

输出：

打印OpSet导入信息: , Version: 17
模型的输入信息:
Name: images, Type: tensor_type {
  elem_type: 1
  shape {
    dim {
      dim_value: 1
    }
    dim {
      dim_value: 3
    }
    dim {
      dim_value: 640
    }
    dim {
      dim_value: 640
    }
  }
}


模型的输出信息:
Name: output0, Type: tensor_type {
  elem_type: 1
  shape {
    dim {
      dim_value: 1
    }
    dim {
      dim_value: 10
    }
    dim {
      dim_value: 8400
    }
  }
}

输入结构：[1，3，640，640]

输出结构：[1，84，8400]

参考文献

• https://github.com/onnx/onnx
• https://github.com/microsoft/onnxruntime
• https://netron.app/