Mini-Infer (15): OnnxParser 架构设计

1. 为什么是 ONNX？

ONNX (Open Neural Network Exchange) 是目前 AI 行业的“通用语”。PyTorch, TensorFlow, Keras 等所有主流训练框架都能导出为 ONNX。

ONNX 文件的本质是一个 Protocol Buffers (Protobuf) 序列化对象。

  • 文件结构：ModelProto -> GraphProto -> NodeProto (算子) / TensorProto (权重)。

我们的任务就是编写一个“翻译器”，将 ONNX 的这些 Proto 对象，翻译成 Mini-Infer 的 Graph 和 Node 对象。

2. 架构蓝图：模仿 TensorRT

TensorRT 的 ONNX Parser 架构非常优秀，我们将借鉴它的设计思想：注册机制 (Registration) 与 导入器 (Importer)。

我们不希望写一个巨大的 switch-case 来处理所有 ONNX 算子。我们希望每一个 ONNX 算子（如 Conv, Relu）都有一个独立的 Importer 函数。

组件设计

  1. OnnxParser (顶层入口): 负责读取文件、解析 Protobuf 二进制数据，是用户调用的 API。
  2. ModelImporter (模型导入器): 负责遍历 ONNX 图 (GraphProto)，管理 Tensor 映射表（ONNX 里的名字 -> Mini-Infer 里的 Tensor 指针）。
  3. OperatorImporter (算子导入器): 定义了如何将一个特定的 ONNX 算子转换为 Mini-Infer 算子。
  4. OperatorRegistry (注册表): 一个 Map<string, ImporterFunc>，负责将 ONNX OpType (如 “Conv”) 映射到对应的 OperatorImporter。

3. 实现 OnnxParser (onnx_parser.h & .cpp)

OnnxParser 类非常简洁，它主要处理 I/O 和 Protobuf 解析。

A. 依赖管理 (#ifdef)

注意我们在 .cpp 中使用了 #ifdef MINI_INFER_ONNX_ENABLED。这是一个很好的工程实践。因为 Protobuf 是一个很重的依赖，用户可能希望编译一个“精简版”的 Mini-Infer（只运行，不加载 ONNX）。这种宏开关允许用户在 CMake 中灵活配置。

B. 解析流程

  1. 读取文件：使用 std::ifstream 将 .onnx 文件读入 std::vector<char> 缓冲区。
  2. 解析 Protobuf：调用 onnx::ModelProto::ParseFromArray。这一步会将二进制数据反序列化为 C++ 对象树。
  3. 调用 Importer：将解析好的 ModelProto 对象交给 ModelImporter 处理。
1
2
3
4
5
6
7
8
9
10
11
12
13
14
// mini_infer/importers/onnx_parser.cpp

std::unique_ptr<graph::Graph> OnnxParser::parse_from_buffer(const void* buffer, size_t size) {
    // 1. Parse Protobuf
    onnx::ModelProto model;
    if (!model.ParseFromArray(buffer, static_cast<int>(size))) {
        set_error("Failed to parse ONNX protobuf");
        return nullptr;
    }
    
    // 2. Import to Mini-Infer Graph
    auto graph = model_importer_->import_model(model);
    return graph;
}

4. 关键依赖：Protobuf

要让这段代码跑起来，你需要在你的项目中集成 Protobuf 和 ONNX 的 .proto 定义文件。

通常的做法是：

  1. 下载 onnx.proto 文件（来自 ONNX GitHub）。
  2. 在 CMake 构建过程中，使用 protoc 编译器将 onnx.proto 编译为 onnx.pb.h 和 onnx.pb.cc。
  3. 将生成的文件链接到你的 Mini-Infer 库中。

这部分通常在 CMakeLists.txt 中处理，是工程配置的难点，但代码逻辑本身很清晰。

5. 总结与展望

本篇我们搭建了 OnnxParser 的“外壳”。我们有了读取文件和解析 Protobuf 的能力。

但是，ModelImporter 和 OperatorRegistry 目前还只是空壳（Forward Declarations）。引擎还不知道如何处理具体的算子。