3. ONNX模型基础:理解ONNX格式、模型结构与张量概念
3.1 为什么Java开发者需要理解ONNX内部结构
在Java服务端直接调用ONNX Runtime执行推理时,多数报错信息都指向模型内部结构。例如,常见的报错 Invalid Node: Node (Conv_3) has input size 2 not in [3, 3],如果不理解ONNX中算子(Operator)的输入输出规则,很难定位是模型导出问题还是调用参数问题。另一个高频场景是:PyTorch导出的模型在Python端运行正常,但Java端加载时提示 Unsupported data type: INT8,这需要你理解ONNX张量的数据类型编码规则。因此,掌握ONNX格式的底层结构,是Java部署排障的第一道基本功。
3.2 ONNX文件格式:Protobuf与序列化
ONNX模型文件(.onnx)使用Google Protobuf(Protocol Buffers)进行序列化。Protobuf是一种二进制序列化协议,比JSON/XML更紧凑,解析速度更快。ONNX的Protobuf定义文件(onnx.proto)规定了模型的所有数据结构。
一个完整的ONNX文件在磁盘上的逻辑结构如下:
ModelProto
├── ir_version: 8 # ONNX IR版本号
├── opset_import: [domain: "", version: 17] # 算子集版本
├── producer_name: "pytorch" # 导出工具
├── producer_version: "2.1.0"
├── graph: GraphProto # 核心计算图
│ ├── node: [NodeProto, NodeProto, ...] # 算子节点列表
│ ├── initializer: [TensorProto, ...] # 初始权重(常量)
│ ├── input: [ValueInfoProto, ...] # 图输入
│ ├── output: [ValueInfoProto, ...] # 图输出
│ └── value_info: [...] # 中间张量形状信息
└── metadata_props: [...] # 附加元数据
其中,ir_version 表示ONNX中间表示(IR)的版本,不同版本支持的算子特性和属性不同。opset_import 声明了该模型使用的算子集版本,Java端ONNX Runtime必须支持该版本才能正确加载。如果Java端报错 Unsupported opset version,通常就是这里不匹配。
3.3 计算图结构:节点、边与张量
ONNX模型的核心是计算图(Graph)。计算图由三类基本元素构成:
- 节点(Node):表示一个数学操作,如卷积(Conv)、矩阵乘(MatMul)、激活(Relu)等。每个节点有
op_type(算子类型)、inputs(输入张量名列表)、outputs(输出张量名列表)和attribute(算子属性,如卷积的kernel_shape、strides)。 - 张量(Tensor):数据的基本载体,具有名称、数据类型、形状(shape)和原始数据。张量在图中以“边”的形式连接节点。
- 边(Edge):张量名作为边的标识,连接生产者和消费者节点。例如,节点A的输出张量名
conv1_out如果被节点B作为输入引用,则A→B之间存在一条边。
下面是一个典型的ONNX计算图片段,展示了一个两层的全连接网络结构:
graph: {
node: {
input: ["input", "fc1_weight", "fc1_bias"]
output: ["fc1_out"]
op_type: "MatMul" // 实际中常为Gemm
}
node: {
input: ["fc1_out"]
output: ["relu1_out"]
op_type: "Relu"
}
node: {
input: ["relu1_out", "fc2_weight", "fc2_bias"]
output: ["output"]
op_type: "MatMul"
}
initializer: [
{ name: "fc1_weight", data_type: FLOAT, dims: [784, 512] },
{ name: "fc1_bias", data_type: FLOAT, dims: [512] },
{ name: "fc2_weight", data_type: FLOAT, dims: [512, 10] },
{ name: "fc2_bias", data_type: FLOAT, dims: [10] }
]
input: [{ name: "input", type: { tensor_type: { elem_type: FLOAT, shape: [1, 784] }}}]
output: [{ name: "output", type: { tensor_type: { elem_type: FLOAT, shape: [1, 10] }}}]
}
注意,initializer 中的权重(如 fc1_weight)是常量张量,它们不来自任何节点,而是直接嵌入模型文件。Java端加载时,这些权重会被反序列化并存储在内存中,供后续算子调用。
input 列表只包含动态输入(即推理时每次需要传入的数据),而 initializer 中的权重虽然也是节点输入,但它们是静态常量,不会在每次推理时重新传入。在Java端调用时,你只需要为 input 列表中的张量提供数据。
3.4 张量概念:形状、数据类型与内存布局
张量(Tensor)是ONNX中数据的基本单元,可以理解为多维数组。理解张量需要掌握以下四个维度:
| 属性 | 说明 | 示例 |
|---|---|---|
| 形状(Shape) | 每个维度的大小,用整数数组表示 | [1, 3, 224, 224] 表示1张3通道224×224图片 |
| 数据类型(Data Type) | 每个元素的存储类型 | FLOAT(32位浮点)、INT64、DOUBLE |
| 元素总数 | 所有维度大小的乘积 | 1×3×224×224 = 150528 |
| 内存布局 | 数据在内存中的排列顺序(行主序) | 按行优先连续存储 |
ONNX中常见的数据类型枚举值如下(Java端可通过 OnnxTensor.createTensor 时指定):
// ONNX TensorProto.DataType 枚举(部分)
FLOAT = 1 // 32位浮点
UINT8 = 2 // 8位无符号整数
INT8 = 3 // 8位有符号整数
UINT16 = 4
INT16 = 5
INT32 = 6
INT64 = 7 // 64位整数,常用于shape、索引
DOUBLE = 11 // 64位浮点
BOOL = 9
在Java端,当你需要将输入数据转换为ONNX张量时,必须确保Java数组的类型与ONNX张量的 elem_type 匹配。例如,如果模型输入是 FLOAT 类型,你应使用 float[] 或 float[][] 数组,而不是 double[],否则会抛出类型不匹配异常。
3.5 可视化案例:一个实际模型的图结构
下面以一个简化的图像分类模型(输入1×3×32×32,输出1×10)为例,展示其ONNX计算图的可视化结构。该图展示了数据从输入到输出的流动路径,以及各节点的连接关系。
上图展示了几个关键特征:initializer(权重)以虚线箭头指向对应的Conv节点,表示它们是常量输入;中间张量(如 conv1_out)在图中以绿色/橙色小标签标注,它们不占用额外存储,只是逻辑上的数据流标识。在Java端调试时,你可以通过ONNX Runtime的 run() 方法指定输出中间张量名来获取这些值,用于验证计算正确性。
3.6 方法流程:Java端如何解析和加载ONNX模型
在Java中,使用ONNX Runtime加载模型的典型流程如下:
- 创建OrtSession:通过
OrtEnvironment.getEnvironment()获取全局环境,然后调用session = env.createSession(modelPath, sessionOptions)加载模型。 - 检查输入输出信息:通过
session.getInputInfo()和session.getOutputInfo()获取输入输出张量的名称、形状和数据类型。这一步可以验证你对模型结构的理解是否正确。 - 准备输入张量:根据输入信息创建
OnnxTensor。例如,如果输入是FLOAT类型且形状为[1, 3, 224, 224],你需要将图像数据转换为float[]并调用OnnxTensor.createTensor(env, floatArray, shape)。 - 执行推理:调用
session.run(Collections.singletonMap("input", inputTensor)),返回结果中包含输出张量。 - 解析输出:从
OnnxTensor中提取Java数组,例如result.get(0).getFloatBuffer()获取浮点输出。
// Java端加载ONNX模型的核心代码片段
OrtEnvironment env = OrtEnvironment.getEnvironment();
OrtSession session = env.createSession("model.onnx", new OrtSession.SessionOptions());
// 查看输入信息
Map<String, NodeInfo> inputInfo = session.getInputInfo();
for (Map.Entry<String, NodeInfo> entry : inputInfo.entrySet()) {
System.out.println("输入: " + entry.getKey() +
" 形状: " + entry.getValue().getInfo().getShape() +
" 类型: " + entry.getValue().getInfo().getType());
}
// 创建输入张量(假设输入为 [1, 3, 224, 224] 的FLOAT)
float[] imageData = new float[1 * 3 * 224 * 224]; // 填充图像数据
long[] shape = {1, 3, 224, 224};
OnnxTensor inputTensor = OnnxTensor.createTensor(env, imageData, shape);
// 执行推理
try (OrtSession.Result result = session.run(Collections.singletonMap("input", inputTensor))) {
OnnxTensor outputTensor = result.get(0).getValue();
float[] outputData = outputTensor.getFloatBuffer().array();
// 处理输出数据
}
3.7 常见误区与复盘
很多初学者在Java端调用时,试图将模型的权重(initializer)也作为输入传入。实际上,权重已经内嵌在模型中,你只需要为图的
input 列表中的张量提供数据。如果传入多余的输入,ONNX Runtime会报错 Invalid Argument: Unexpected input name。
PyTorch默认使用
torch.float32,导出ONNX后数据类型为 FLOAT。但如果你在Java端使用 double[] 数组创建张量,会得到 DOUBLE 类型,与模型要求的 FLOAT 不匹配,导致 Type Mismatch 异常。务必使用 float[]。
ONNX中图像张量的标准布局是
[N, C, H, W](批大小、通道数、高度、宽度)。如果你在Java端将图像数据按 [H, W, C] 排列(如某些图像库的默认布局),会导致推理结果完全错误。需要先进行维度转换(permute),再填充到 float[] 中。
session.getInputInfo() 完全一致。最后,可以尝试在ONNX Runtime中开启 SessionOptions.setOptimizationLevel(ORT_ENABLE_BASIC) 来排除图优化导致的数值差异。