3. ONNX模型基础
3.1 ONNX格式的本质:不止是一个“文件”
在接触OpenVINO部署时,很多开发者习惯把ONNX看作“一种模型文件”,就像.pt或.h5一样。但ONNX(Open Neural Network Exchange)本质上是一个计算图中间表示(Intermediate Representation, IR),它定义了一套与框架无关的算子集合和图描述协议。ONNX文件(通常为.onnx,采用Protobuf序列化)保存的是模型的静态计算图——即数据如何从输入张量经过一系列算子流动到输出张量。
理解这一点对后续OpenVINO部署至关重要:OpenVINO的模型优化器(MO)并不“运行”ONNX模型,而是解析ONNX图结构,将其映射为中间表示(IR),再经过编译生成推理引擎可执行的插件格式。因此,ONNX图的拓扑结构、算子类型、张量形状、属性值是否规范,直接决定了MO能否成功转换以及转换后推理性能的优劣。
3.2 ONNX模型结构:从Proto到Graph
一个ONNX模型在逻辑上由以下层级组成:
- ModelProto:最外层容器,包含
ir_version(IR版本)、opset_import(算子集版本)、producer_name(导出工具)等元信息,以及核心的graph字段。 - GraphProto:描述计算图,包含
node(算子节点列表)、initializer(初始权重张量)、input(图输入)、output(图输出)、value_info(中间张量形状推断结果)。 - NodeProto:单个算子节点,包含
op_type(算子类型,如Conv、Relu)、input(输入张量名列表)、output(输出张量名列表)、attribute(算子属性,如卷积的kernel_shape、strides)。 - TensorProto:用于存放
initializer中的权重数据,以多维数组形式存储,支持float32、int64等多种数据类型。
下图展示了ONNX模型的结构层级与数据流关系:
图中展示了ONNX模型的核心结构:ModelProto包裹GraphProto,后者包含输入输出定义、权重初始化器以及有序的算子节点列表。值得注意,ONNX图是静态的——所有控制流(如if-else、循环)都必须显式转换为ONNX的If、Loop算子,而不能保留Python级别的逻辑分支。
3.3 导出ONNX的方法与流程
主流深度学习框架均提供了ONNX导出接口。以PyTorch为例,最常用的是torch.onnx.export。其核心流程如下:
- 准备模型与输入:将模型切换到
eval()模式,并构造一个与训练时输入形状一致的示例张量(dummy input)。 - 调用导出API:指定输出文件路径、输入输出名称、动态轴(如有)、算子集版本等参数。
- 验证导出结果:使用
onnx.checker检查模型合法性,并用onnxruntime跑一遍推理,对比与原框架输出的数值误差。
import torch
import torchvision.models as models
model = models.resnet18(pretrained=True)
model.eval()
dummy_input = torch.randn(1, 3, 224, 224)
torch.onnx.export(
model, # 要导出的模型
dummy_input, # 示例输入
"resnet18.onnx", # 输出文件路径
input_names=["input"], # 输入张量名
output_names=["output"], # 输出张量名
dynamic_axes={"input": {0: "batch_size"}, # 动态batch
"output": {0: "batch_size"}},
opset_version=17, # 算子集版本
do_constant_folding=True # 常量折叠优化
)
导出后,建议立即执行以下验证:
import onnx
import onnxruntime as ort
import numpy as np
# 1. 结构检查
onnx_model = onnx.load("resnet18.onnx")
onnx.checker.check_model(onnx_model)
# 2. 推理对比
ort_session = ort.InferenceSession("resnet18.onnx")
ort_output = ort_session.run(None, {"input": dummy_input.numpy()})[0]
torch_output = model(dummy_input).detach().numpy()
print("最大误差:", np.abs(ort_output - torch_output).max())
3.4 导出ONNX的常见误区与注意事项
| 误区/风险点 | 可能后果 | 正确做法 |
|---|---|---|
忘记调用model.eval() | BN层使用batch统计量,推理结果错误 | 导出前必须切换到eval模式 |
| 使用动态shape但未声明dynamic_axes | 导出的模型固定输入尺寸,部署时无法适配不同分辨率 | 在dynamic_axes中明确声明动态维度 |
| 模型中包含Python原生控制流(if/for) | 导出失败或图被静态展开,失去动态逻辑 | 改用torch.where、torch.max等张量操作替代 |
| opset版本设置过低 | 新算子(如Attention)无法导出,或转换OpenVINO时算子不支持 | 根据算子需求选择opset≥13,推荐17或更高 |
| 未做常量折叠 | 图中残留大量Shape/Gather子图,增加转换负担 | 设置do_constant_folding=True |
| 自定义算子未注册 | 导出报错或生成不规范的custom op | 使用torch.onnx.register_custom_op_symbolic注册映射 |
aten::unique或aten::nonzero算子,在ONNX中可能映射为不规范的子图,导致MO转换失败。因此,导出后务必用onnxruntime验证数值一致性,再用mo --input_model model.onnx试转换,尽早暴露问题。
3.5 案例复盘:一个典型的导出失败与修复
场景:某检测模型使用torchvision.ops.nms(非极大值抑制)进行后处理,直接导出ONNX时抛出RuntimeError: Exporting the operator 'nms' to ONNX opset version 17 is not supported.
问题分析:NMS算子包含循环和条件判断,ONNX标准算子集中没有直接对应的算子。PyTorch的ONNX导出器无法将torchvision.ops.nms映射为ONNX原生算子。
解决方案:将NMS后处理从模型中剥离,只导出主干网络部分;NMS逻辑在部署端用OpenVINO的NonMaxSuppression算子或外部代码实现。修改后的导出代码:
# 只导出不含NMS的检测头
class DetectionHead(nn.Module):
def forward(self, x):
# 回归分支、分类分支...
return boxes, scores
torch.onnx.export(
DetectionHead(), dummy_input,
"det_head.onnx",
opset_version=17,
input_names=["input"],
output_names=["boxes", "scores"]
)
在OpenVINO部署时,使用openvino.runtime.opset9.nms或opset9.non_max_suppression构建后处理,或直接使用OpenVINO的CompiledModel配合CPU插件自带的NMS扩展。