从零构建轻量级推理引擎 OInfer（二）：ONNX 模型解析与计算图构建

本篇深入剖析 OInfer 如何从 ONNX 模型文件（.onnx）解析出完整的计算图，涵盖 Protobuf 反序列化、权重加载、节点创建、Kahn 拓扑排序算法与形状推断注册表机制。

---

1. ONNX 格式概述

ONNX（Open Neural Network Exchange）是微软和 Facebook 联合发起的开放神经网络交换格式，旨在实现不同深度学习框架之间的模型互操作。理解 ONNX 的数据组织方式，是构建推理引擎的第一步。

ONNX 使用 Google 的 Protocol Buffers (Protobuf) 作为序列化格式。Protobuf 是一种语言无关、平台无关的结构化数据序列化协议，相比 JSON 或 XML，其二进制格式更紧凑、解析速度更快，非常适合存储包含大量权重数据的模型文件。OInfer 通过 protoc 编译器将 onnx.proto 编译为 C++ 代码（onnx.pb.h/cc），生成的类提供类型安全的访问接口。

ONNX 模型文件的核心数据层次结构如下：

ModelProto                        ← 模型顶层容器
 ├── ir_version                   ← ONNX IR 版本号
 ├── opset_import[]               ← 算子集版本
 └── GraphProto                   ← 计算图（核心）
      ├── initializer[]           ─── 权重/偏置数据 (TensorProto)
      ├── input[]                 ─── 输入描述 (ValueInfoProto)
      ├── output[]                ─── 输出描述 (ValueInfoProto)
      ├── node[]                  ─── 算子节点 (NodeProto)
      └── value_info[]            ─── 中间张量描述 (ValueInfoProto)

每个 NodeProto 描述一个计算操作，包含：

• op_type：算子类型标识符（如 "Conv", "Relu", "MaxPool"）
• input[]：输入张量名称列表（字符串数组）
• output[]：输出张量名称列表
• attribute[]：算子属性（如卷积的 pads, strides, dilations 等）

关键设计哲学 ：ONNX 中张量不直接内联在节点中，而是通过 名称字符串 构成引用关系。一个节点的输出张量名称可以出现在另一个节点的输入列表中，由此形成数据流图。这意味着解析器需要维护一个 全局名称→张量信息 的映射表，将名称引用解析为实际的指针关系。OInfer 的 Graph::tensors（unordered_map<string, unique_ptr<TensorInfo>>）正是承担这一职责的数据结构。

---

2. OInfer 的内部图表示

OInfer 定义了三个核心数据结构来表示经过解析的计算图。它们之间通过指针建立联系，形成一个高效的内存内图表示。

2.1 TensorInfo：张量元信息

enum class TensorRole : int8_t {
    INPUT        = 0,  // 模型输入（运行时由外部设置）
    OUTPUT       = 1,  // 模型输出
    CONSTANT     = 2,  // 权重/偏置（来自 initializer，数据已加载）
    INTERMEDIATE = 3,  // 中间激活值（运行时动态分配）
};

struct TensorInfo {
    std::string          name;
    std::vector<int64_t> shape;
    int32_t              dtype = 1;  // ONNX DataType, 1 = FLOAT32
    TensorRole           role  = TensorRole::INTERMEDIATE;
    void*                data       = nullptr;  // CONSTANT 时持有权重数据
    size_t               data_bytes = 0;
};

TensorRole 枚举是理解后续资源分配策略的关键。不同角色的张量在 Backend 层有完全不同的处理方式：

┌──────────────────────────────────────────────────────────┐
│                    TensorRole 决策树                      │
├──────────────┬──────────────┬───────────────┬────────────┤
│    INPUT     │   CONSTANT   │ INTERMEDIATE  │   OUTPUT   │
├──────────────┼──────────────┼───────────────┼────────────┤
│ 外部传入数据   │ 加载权重数据   │ 运行时分配空间   │ 读取结果    │
│ NCHW→NHWC    │ OIHW→OHWI    │ 按推断shape    │ 返回给用户   │
│ 运行时填充     │ 加载后不变     │ 自动管理       │ 运行后读取   │
└──────────────┴──────────────┴───────────────┴────────────┘

INPUT 张量的数据在推理时由用户通过 SetInput() 方法填入（需要做 NCHW→NHWC 布局转换）；CONSTANT 张量（权重、偏置）在模型加载时一次性从文件中读取并转换布局（OIHW→OHWI）；INTERMEDIATE 张量的形状在形状推断阶段确定，内存在 Backend 初始化时分配；OUTPUT 张量在推理完成后由用户通过 GetOutput() 方法读取。

2.2 NodeInfo：算子节点

struct NodeInfo {
    std::string name;
    std::string op_type;
    std::vector<TensorInfo*> inputs;   // 指向 Graph::tensors 中的元素
    std::vector<TensorInfo*> outputs;
    std::unordered_map<std::string, onnx::AttributeProto> attrs;
};

这里有一个重要的设计选择：inputs 和 outputs 存储的是 裸指针 而非智能指针或名称字符串。这些指针直接指向 Graph::tensors 容器中的 TensorInfo 对象。这种设计带来两个好处：一是避免了名称查找的 O(1) 哈希开销（在形状推断等高频访问路径上累积可观）；二是当形状推断函数写入某个 TensorInfo 的 shape 字段后，所有引用该张量的下游节点可以立即通过指针读取到最新的形状信息，无需额外的通知机制。

属性使用 onnx::AttributeProto 原样存储，保持了与 ONNX 规范的完整兼容性，各算子的形状推断和执行逻辑可以根据需要从中提取 ints、floats、strings 等不同类型的属性值。

2.3 Graph：计算图

struct Graph {
    std::unordered_map<std::string, std::unique_ptr<TensorInfo>> tensors;
    std::vector<std::unique_ptr<NodeInfo>>                       nodes;
    std::vector<NodeInfo*>                                       topo_order;
};

tensors 拥有所有 TensorInfo 的所有权（通过 unique_ptr），nodes 拥有所有 NodeInfo 的所有权，而 topo_order 是一个指针向量，指向 nodes 中的元素，表示经过拓扑排序后的执行顺序。这种"所有权在容器，引用在 topo_order"的分离设计，使得排序操作不需要移动任何数据。

---

3. 图构建五阶段流水线

BuildGraphFromOnnx() 是整个图构建流程的入口函数。它先从文件读取并反序列化 ModelProto，然后按严格顺序执行五个加载阶段，每个阶段处理 GraphProto 的一部分数据：

  onnx::GraphProto
        │
        ├─── ① loadInitializers()   加载权重 → TensorRole::CONSTANT
        │
        ├─── ② loadInputs()         加载输入描述 → TensorRole::INPUT
        │
        ├─── ③ createNodes()        创建算子节点，建立 tensor 引用
        │
        ├─── ④ loadOutputs()        标记输出 → TensorRole::OUTPUT
        │
        └─── ⑤ loadValueInfo()      补充中间张量形状信息
              │
              ▼
        topoSort() + inferShapes()

阶段顺序不可随意调整 ------这是一个经常被忽略但至关重要的细节。ONNX 规范中，initializer 中出现的张量名称 同时也会出现在 input 列表中 （这是 ONNX 的历史设计，为了向后兼容）。因此 loadInitializers() 必须先于 loadInputs() 执行，确保权重张量先被标记为 CONSTANT。当 loadInputs() 再次遇到同名张量时，通过角色检查跳过覆盖：

static void loadInputs(const onnx::GraphProto& g, Graph& graph) {
    for (const auto& v : g.input()) {
        TensorInfo* t = getOrCreate(graph, v.name());
        // 已被 loadInitializers() 标记为 CONSTANT 的权重张量不覆盖
        if (t->role != TensorRole::CONSTANT) {
            t->role = TensorRole::INPUT;
        }
        // ... 填充 shape / dtype
    }
}

如果颠倒顺序，权重张量会被错误地标记为 INPUT，导致 Backend 层尝试为其分配可修改的激活内存而非加载权重数据。

权重数据加载

loadInitializers() 从 TensorProto 中提取权重数据。ONNX 支持两种数据存储格式，OInfer 依次尝试：

TensorProto
  ├── raw_data      ─── 原始二进制字节流（紧凑，PyTorch 导出常用，优先）
  └── float_data[]  ─── Protobuf repeated float 字段（可读性好，fallback）

if (init.has_raw_data()) {
    std::memcpy(t->data, init.raw_data().data(), bytes);
} else if (init.float_data_size() > 0) {
    std::memcpy(t->data, init.float_data().data(), bytes);
}

数据通过 malloc 分配并直接 memcpy，在 TensorInfo 的析构函数中 free。当前实现仅支持 float32 类型（dtype == 1），这覆盖了绝大多数常见模型场景。

getOrCreate 辅助函数

五个加载阶段都使用同一个辅助函数来访问或创建张量：

static TensorInfo* getOrCreate(Graph& graph, const std::string& name) {
    auto& ptr = graph.tensors[name];
    if (!ptr) {
        ptr       = std::make_unique<TensorInfo>();
        ptr->name = name;
    }
    return ptr.get();
}

这个函数利用 unordered_map 的 operator[] 的"不存在则插入默认值"语义，用一行代码完成了"查找-存在则返回-不存在则创建"的逻辑。返回裸指针，所有权始终留在 graph.tensors 容器中。

---

4. 拓扑排序：Kahn 算法

计算图中存在数据依赖关系。如果 Conv 的输出张量是 Relu 的输入张量，那么 Conv 必须先于 Relu 执行。拓扑排序的目的就是确定一个合法的执行顺序，使得每个节点执行时，其所有输入数据都已就绪。

OInfer 使用 Kahn 算法（基于 BFS 的拓扑排序）来实现。相比基于 DFS 的拓扑排序，Kahn 算法的优势在于能自然地检测环（如果排序结果的节点数少于总节点数，则存在环）。

算法详解

                      ┌────────────────────────────────────┐
                      │ Step 1: 初始 available 集合         │
                      │   = 所有 INPUT + CONSTANT 张量      │
                      │   （这些张量不依赖任何节点输出）     │
                      └─────────────────┬──────────────────┘
                                        │
                ┌───────────────────────▼───────────────────────┐
                │ Step 2: 计算每个节点的 pending 计数            │
                │   pending[node] = 该节点输入中不在 available   │
                │                   集合中的张量数               │
                │   (即"尚未就绪的依赖数")                      │
                └───────────────────────┬───────────────────────┘
                                        │
                    ┌───────────────────▼──────────────────┐
                    │ Step 3: 建立 tensor→消费者 的映射     │
                    │   consumers[tensor_name] = [nodes..] │
                    └───────────────────┬──────────────────┘
                                        │
                    ┌───────────────────▼──────────────────┐
                    │ Step 4: 将 pending == 0 的节点入队    │
                    └───────────────────┬──────────────────┘
                                        │
                    ┌───────────────────▼──────────────────┐
                    │  while ready 非空:                    │
                    │    cur = ready.dequeue()              │
                    │    topo_order.push_back(cur)          │
                    │    for out in cur->outputs:           │
                    │      available.insert(out)            │
                    │      for consumer in consumers[out]:  │
                    │        if --pending[consumer] == 0:   │
                    │          ready.enqueue(consumer)      │
                    └──────────────────────────────────────┘

算法的核心思想是：一个节点只有在其所有输入张量都已"可用"时才能被执行。初始时，模型的输入和权重天然可用；每当一个节点执行完毕，其输出张量变为可用，可能"解锁"下游节点。

实例演示

以 conv_relu_model.onnx 为例，这是一个包含 Conv 和 Relu 两个算子的简单模型：

input ──┐
        ├──▶ Conv ──▶ conv_out ──▶ Relu ──▶ output
weight ─┘      ▲
bias ──────────┘

Step 1: available = {input, weight, bias}

Step 2: 
  Conv: pending = 0  (input ✓, weight ✓, bias ✓ → 全部可用)
  Relu: pending = 1  (conv_out ✗ → 尚不可用)

Step 3: consumers = {"conv_out": [Relu]}

Step 4: ready = [Conv]

迭代:
  弹出 Conv → topo_order = [Conv]
    available += {conv_out}
    Relu: --pending = 0 → 入队
  弹出 Relu → topo_order = [Conv, Relu]

最终: topo_order = [Conv, Relu]  ✓

OInfer 还加入了完整性检查：如果排序后 topo_order.size() != nodes.size()，说明图中存在环（这不应该出现在合法的 ONNX 模型中），引擎输出警告信息。

---

5. 形状推断机制

形状推断（Shape Inference）的目标是：在模型加载阶段就确定所有中间张量的精确 shape，以便后续 Backend 层为每个张量精确分配内存。如果缺少形状推断，Backend 就不知道应该为中间张量分配多大的内存空间。

5.1 注册表模式

OInfer 采用 单例注册表 + 静态自动注册 的设计模式。这种模式在 C++ 项目中非常经典，其核心思想是：每个算子的形状推断逻辑在自己的编译单元（.cpp 文件）中注册，全局注册表在运行时通过查表的方式调用对应的推断函数。

┌──────────────────────────────────────────────┐
│           ShapeInferRegistry (单例)            │
│                                              │
│  fns_: unordered_map<string, ShapeInferFn>   │
│                                              │
│  "Conv"  → [lambda: Conv 形状推断逻辑]        │
│  "Relu"  → [lambda: Relu 形状推断逻辑]        │
└──────────────────────────────────────────────┘
          ▲                       ▲
          │                       │
    conv_kernel.cpp         relu_kernel.cpp
    REGISTER_SHAPE_INFER    REGISTER_SHAPE_INFER

宏定义展开后创建一个文件作用域的静态全局变量 ShapeInferRegistrar，其构造函数在程序启动（main() 之前）自动调用：

#define REGISTER_SHAPE_INFER(op, ...) \
    static ShapeInferRegistrar _shape_infer_reg_##__LINE__{ (op), (__VA_ARGS__) }

ShapeInferRegistrar::ShapeInferRegistrar(const std::string& op_type, ShapeInferFn fn) {
    ShapeInferRegistry::Instance().Register(op_type, std::move(fn));
}

ShapeInferRegistry::Instance() 使用 Meyer's Singleton 模式------C++11 标准保证函数内的 static 变量在首次进入时线程安全地初始化。这避免了 C++ 中臭名昭著的 SIOF（Static Initialization Order Fiasco） 问题：即使 ShapeInferRegistrar 的构造函数在另一个编译单元的全局变量初始化中被调用，注册表单例也一定已经准备就绪。

5.2 Conv 形状推断的实现

Conv 算子的输出形状计算是标准的卷积维度公式。OInfer 支持 strides、pads、dilations 三个属性，均可从节点的 attrs 中提取（若缺失则使用默认值）：

out_h = ⌊(in_h + pad_top + pad_bottom - dilation_h × (kh - 1) - 1) / stride_h⌋ + 1
out_w = ⌊(in_w + pad_left + pad_right  - dilation_w × (kw - 1) - 1) / stride_w⌋ + 1

输入: input.shape  = [N, C_in, H_in, W_in]     (ONNX NCHW)
      weight.shape = [C_out, C_in, kH, kW]      (ONNX OIHW)

输出: output.shape = [N, C_out, H_out, W_out]   (ONNX NCHW)

注意这里的推断在 ONNX 原始的 NCHW 坐标系中进行，到 NHWC 的转换推迟到 Backend 层。这种"推断用原始布局，执行用内部布局"的分层策略保持了 Graph 层的纯粹性------它只关心拓扑和形状，不涉及任何执行层面的布局优化。

5.3 推断执行流程

形状推断按拓扑顺序逐节点执行，这保证了处理某节点时其所有输入的 shape 已经就绪：

static void inferShapes(Graph& graph) {
    for (NodeInfo* node : graph.topo_order) {
        // 1. 检查输出 shape 是否需要推断（可能模型已经提供）
        // 2. 检查所有输入 shape 是否就绪
        // 3. 查注册表执行推断
        inferNodeShape(node);
    }
}

Fallback 策略：对于未注册形状推断函数的算子（如未来可能扩展的 BatchNorm、Add 等），系统提供了一个通用的 fallback------将第一个输入的 shape 直接复制到所有输出。这对 element-wise 算子（Relu, Sigmoid, Tanh 等）是正确的，因为它们不改变张量形状。当然，对于 Reshape、Flatten 等会改变形状的算子，必须提供专门的推断函数。

---

6. 完整流程示例

以 conv_relu_model.onnx（Conv + Relu 模型，输入 [1,1,3,3]，卷积核 [1,1,2,2]，padding=1,1,1,1）为例，完整跟踪图构建过程：

Step 1: loadInitializers()
  ├── weight: shape=[1,1,2,2], role=CONSTANT, data=loaded (16 bytes)
  └── bias:   shape=[1],       role=CONSTANT, data=loaded (4 bytes)

Step 2: loadInputs()
  └── input: shape=[1,1,3,3], role=INPUT
      (weight 已是 CONSTANT，跳过覆盖)

Step 3: createNodes()
  ├── Node[0]: op="Conv",  inputs=[input, weight, bias], outputs=[conv_out]
  │            attrs: pads=[1,1,1,1]
  └── Node[1]: op="Relu",  inputs=[conv_out], outputs=[output]

Step 4: loadOutputs()
  └── output: role=OUTPUT

Step 5: loadValueInfo()
  └── conv_out: 补充形状信息（若模型导出时已填写）

Step 6: topoSort()
  └── topo_order = [Conv, Relu]

Step 7: inferShapes()
  ├── Conv: input=[1,1,3,3], weight=[1,1,2,2], pads=[1,1,1,1]
  │         H_out = (3 + 1 + 1 - 1×(2-1) - 1) / 1 + 1 = 4
  │         W_out = (3 + 1 + 1 - 1×(2-1) - 1) / 1 + 1 = 4
  │         → output=[1, 1, 4, 4]
  └── Relu: input=[1,1,4,4] → output=[1,1,4,4]  (直通)

至此，Graph 构建完成。每个 TensorInfo 都有了确定的 shape、role 和数据（如果是 CONSTANT），为后续 Backend 的资源分配提供了全部必要信息。

---

7. 与工业级框架的对比

OInfer 的图构建流程虽然简化，但与工业级框架的核心逻辑高度一致：

特性	OInfer	MNN	NCNN
模型格式	ONNX (Protobuf)	MNN FlatBuffer / ONNX	NCNN 自定义 + ONNX 转换
图表示	Graph{tensors, nodes}	Net{ops, tensors}	Net{layers, blobs}
拓扑排序	Kahn 算法	离线工具预排序	离线工具预排序
形状推断	运行时按拓扑序	运行时 + 离线预推断	运行时
算子注册	静态注册表	工厂模式 + 注册	工厂模式 + 注册

主要差异在于：工业级框架通常在模型转换阶段（离线工具）就完成了拓扑排序和部分形状推断，运行时只需验证；而 OInfer 将所有逻辑放在运行时，虽然有启动开销，但实现更简单直观。

---

8. 小结

本篇完整展示了 OInfer 从 ONNX 文件到内部计算图的构建过程。整个流程可以浓缩为一张表：

阶段	输入	输出	关键技术
Protobuf 解析	.onnx 二进制文件	ModelProto 对象	Protobuf 反序列化
五阶段加载	GraphProto 的各字段	Graph{tensors, nodes}	名称索引、角色标记、顺序约束
拓扑排序	nodes	topo_order	Kahn 算法 (BFS)、环检测
形状推断	topo_order + tensors	完整 shape	单例注册表 + Meyer's Singleton + fallback

图构建完成后，所有信息（拓扑序、精确 shape、权重数据）已经冻结，可以安全地交给 Backend 层进行资源分配和 Pipeline 构建。Graph 对象自此成为只读的"蓝图"。

下一篇将进入 OpenCL 层，详解 Tensor 数据如何映射到 OpenCL Image2D，以及 C4 对齐（Channel-4 Packing）的设计原理与坐标计算。