3、模型解析与加载：ONNX模型格式解析、Protobuf序列化与反序列化、模型元数据读取

好，咱们进入第三章。这一章要啃的，是推理引擎的「入口」——模型解析与加载。

你想想看，无论你的引擎底层优化得多花哨，算子融合得多巧妙，第一步总得先把模型文件读进来吧？ONNX 是目前业界最通用的中间格式，几乎成了事实标准。所以，搞懂 ONNX 怎么存数据、怎么用 Protobuf 把它读出来、怎么提取元数据，是咱们绕不开的硬功夫。

3.1 ONNX 模型格式解析

ONNX 全称 Open Neural Network Exchange，说白了就是「开放神经网络交换格式」。它的设计初衷很朴素：让不同框架（PyTorch、TensorFlow、Caffe2 等）训练出来的模型，能互相串门。

我个人习惯把 ONNX 模型看作一个「计算图」的序列化快照。它用 Protobuf 序列化，存的是二进制数据。但别被「二进制」三个字吓到，它的结构其实很清晰。

3.1.1 ONNX 的文件结构

一个标准的 ONNX 模型文件（.onnx），内部是一个 ModelProto 对象。它长这样：

ModelProto {
    ir_version: int64          // IR 版本号
    opset_import: []OperatorSetIdProto  // 算子集版本
    producer_name: string      // 生成该模型的框架名
    producer_version: string   // 框架版本
    domain: string             // 域名
    model_version: int64       // 模型版本
    doc_string: string         // 文档说明
    graph: GraphProto          // 核心：计算图
    metadata_props: []StringStringEntryProto  // 元数据键值对
}

嗯，这里要注意：graph 才是真正的「肉」。其他字段更像是「包装纸」——告诉你这个模型是谁生的、版本多少、用了哪个算子集。

3.1.2 GraphProto 内部结构

GraphProto 是 ONNX 的灵魂。它包含：

  • node：算子节点列表。每个 NodeProto 包含 op_type（算子类型，如 Conv、Relu）、input 名称列表、output 名称列表、attribute 属性。
  • initializer：常量张量。比如卷积的权重、BN 的均值和方差。这些是训练好的参数，推理时直接加载。
  • input：模型输入信息。包括输入张量的名称、数据类型、形状。
  • output：模型输出信息。结构同 input。
  • value_info：中间张量的形状信息。这个字段经常被忽略，但我在项目中吃过它的亏——有些模型不填 value_info，导致推理时动态 shape 解析困难。

核心要点：ONNX 模型 = 计算图（node 列表）+ 权重数据（initializer）+ 输入输出描述（input/output）。

3.2 Protobuf 序列化与反序列化

ONNX 用 Protobuf（Protocol Buffers）做序列化。Protobuf 是 Google 搞的一种高效数据交换格式，比 XML、JSON 小得多，解析也快得多。

为什么会选 Protobuf？说白了，深度学习模型动不动几百兆，用 JSON 存权重？那解析时间够你喝几杯咖啡的。Protobuf 是二进制格式，体积小、解析快，而且有严格的 schema 约束——你传错字段类型，编译期就报错。

3.2.1 反序列化流程

咱们推理引擎要做的，就是把 .onnx 文件从磁盘读到内存，变成 C++ 对象。流程如下：

  1. 读取文件到 std::string 或 std::vector<char>。
  2. 调用 google::protobuf::io::ArrayInputStream 包装内存数据。
  3. 调用 ModelProto::ParseFromZeroCopyStream() 解析。
  4. 检查解析结果。如果失败，大概率是文件损坏或版本不匹配。

代码示例：

#include <fstream>
#include <google/protobuf/io/zero_copy_stream_impl.h>
#include <onnx/onnx_pb.h>

bool LoadONNXModel(const std::string& path, onnx::ModelProto& model) {
    std::ifstream ifs(path, std::ios::binary | std::ios::ate);
    if (!ifs.is_open()) {
        // 文件打不开，常见于路径错误或权限不足
        return false;
    }

    std::streamsize size = ifs.tellg();
    ifs.seekg(0, std::ios::beg);

    std::vector<char> buffer(size);
    if (!ifs.read(buffer.data(), size)) {
        return false;
    }

    google::protobuf::io::ArrayInputStream ais(buffer.data(), size);
    return model.ParseFromZeroCopyStream(&ais);
}

避坑指南：我曾经在解析大模型（超过 2GB）时踩过坑——直接用 ParseFromIstream 会触发 Protobuf 的 2GB 限制。改用 ParseFromZeroCopyStream + 手动分片读取可以绕过。另外，记得链接 -lprotobuf -lprotoc，别漏了。

3.2.2 序列化（写模型）

虽然推理引擎主要是读模型，但调试时经常需要「写」——比如修改某个节点的属性后重新保存。序列化就是反序列化的逆过程：

bool SaveONNXModel(const onnx::ModelProto& model, const std::string& path) {
    std::ofstream ofs(path, std::ios::binary);
    if (!ofs.is_open()) return false;

    google::protobuf::io::OstreamOutputStream oos(&ofs);
    return model.SerializeToZeroCopyStream(&oos);
}

嗯，这里要注意：序列化时，Protobuf 默认会做一定的压缩（Varint 编码），但不会像 zip 那样压缩权重数据。如果你发现 .onnx 文件比原始权重还大，别慌——那是 Protobuf 的 overhead，通常可以忽略。

3.3 模型元数据读取

元数据是什么？就是「关于数据的数据」。在 ONNX 模型里，元数据包括：

  • 模型名称、版本、作者
  • 输入输出的名称、数据类型、形状
  • 算子集版本（opset version）
  • 自定义的 metadata_props（键值对）

我个人习惯在引擎初始化时，先把元数据全部读出来，存到一个结构体里。这样后面做图优化、内存分配时，不用反复解析 Protobuf 对象——那玩意儿解析一次就够了，反复解析是性能杀手。

3.3.1 读取输入输出信息

输入输出信息藏在 graph.input() 和 graph.output() 里。每个元素是一个 ValueInfoProto，包含：

ValueInfoProto {
    name: string
    doc_string: string
    type: TypeProto {
        tensor_type: TensorTypeProto {
            elem_type: int32  // 数据类型枚举
            shape: TensorShapeProto {
                dim: []TensorShapeProto.Dimension {
                    dim_value: int64   // 固定维度
                    dim_param: string  // 动态维度（如 "batch"）
                }
            }
        }
    }
}

代码示例：

struct ModelIOInfo {
    std::string name;
    int32_t data_type;  // 对应 ONNX 的 TensorProto::DataType
    std::vector<int64_t> shape;
    std::vector<std::string> dynamic_dims;  // 动态维度名称
};

std::vector<ModelIOInfo> ParseInputs(const onnx::GraphProto& graph) {
    std::vector<ModelIOInfo> inputs;
    for (const auto& input : graph.input()) {
        ModelIOInfo info;
        info.name = input.name();

        const auto& type = input.type().tensor_type();
        info.data_type = type.elem_type();

        for (const auto& dim : type.shape().dim()) {
            if (dim.has_dim_value()) {
                info.shape.push_back(dim.dim_value());
                info.dynamic_dims.push_back("");
            } else if (dim.has_dim_param()) {
                info.shape.push_back(-1);  // -1 表示动态
                info.dynamic_dims.push_back(dim.dim_param());
            }
        }
        inputs.push_back(info);
    }
    return inputs;
}

注意：有些模型输入是动态 shape（比如 batch size 可变），dim_value 字段可能不存在，而是用 dim_param 表示。你如果直接读 dim_value() 会得到 0，然后以为形状是 0 维——我刚开始就犯过这个错，排查了半天。

3.3.2 读取算子集版本

算子集版本（opset version）决定了你的引擎支持哪些算子。比如 opset 11 引入了 ScatterND，opset 13 改了 Softmax 的 axis 行为。读取方式：

int GetOpsetVersion(const onnx::ModelProto& model) {
    for (const auto& opset : model.opset_import()) {
        if (opset.domain() == "" || opset.domain() == "ai.onnx") {
            return opset.version();
        }
    }
    return -1;  // 没找到？那这模型可能有问题
}

嗯，这里要注意：domain 为空字符串或 ai.onnx 表示标准 ONNX 算子集。如果有自定义算子，domain 会是其他值（比如 com.nvidia），版本号也要单独处理。

3.3.3 读取 metadata_props

有些模型会在 metadata_props 里塞一些额外信息，比如：

  • 训练时的超参数
  • 输入数据的预处理方式（mean/std）
  • 模型的作者、描述

读取方式很简单：

std::string GetMetadata(const onnx::ModelProto& model, const std::string& key) {
    for (const auto& prop : model.metadata_props()) {
        if (prop.key() == key) {
            return prop.value();
        }
    }
    return "";
}

个人经验：我在对接某个第三方模型时，发现模型推理结果总是不对。后来打印 metadata_props，发现里面有个 input_scale 字段——原来输入需要除以 255 做归一化，而模型本身没做这个预处理。metadata_props 有时候就是「说明书」，不看会出大问题。

3.4 实战：一个完整的模型加载器

把上面这些串起来，一个最小可用的模型加载器大概长这样：

class ONNXModelLoader {
public:
    bool Load(const std::string& path) {
        // 1. 反序列化
        if (!LoadONNXModel(path, model_)) {
            return false;
        }

        // 2. 提取元数据
        ir_version_ = model_.ir_version();
        opset_version_ = GetOpsetVersion(model_);
        producer_ = model_.producer_name();

        // 3. 解析图结构
        const auto& graph = model_.graph();
        inputs_ = ParseInputs(graph);
        outputs_ = ParseOutputs(graph);
        initializers_ = ParseInitializers(graph);
        nodes_ = ParseNodes(graph);

        // 4. 做一些基本校验
        if (nodes_.empty()) {
            // 空图？这模型怕不是个空壳
            return false;
        }

        return true;
    }

    // ... 各种 getter 方法

private:
    onnx::ModelProto model_;
    int ir_version_;
    int opset_version_;
    std::string producer_;
    std::vector<ModelIOInfo> inputs_;
    std::vector<ModelIOInfo> outputs_;
    std::vector<TensorProto> initializers_;
    std::vector<NodeProto> nodes_;
};

这个加载器虽然简陋，但已经能应付 80% 的 ONNX 模型了。后面章节我们会逐步给它加上图优化、内存预分配、算子注册等功能——那才是推理引擎真正发力的地方。

3.5 本章小结

咱们这一章干了三件事：

  1. 解析 ONNX 格式：搞懂了 ModelProto → GraphProto → NodeProto 的层级关系。
  2. 玩转 Protobuf：学会了反序列化（读模型）和序列化（写模型），顺便避了个 2GB 大文件的坑。
  3. 提取元数据：把输入输出信息、算子集版本、自定义属性都读了出来，为后续的图优化做准备。

下一章，咱们要开始真正「动刀」了——对计算图做优化。到时候你会发现，模型解析只是开胃菜，图优化才是硬骨头。不过别急，先把这一章的基础打牢，后面才能跑得快。

嗯，今天就到这儿。有问题随时找我。

← 上一章 📖 返回目录 下一章 →