前言
将模型转为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/