第4章前端解析与模型导入

第四章：前端解析与模型导入

好，咱们今天聊聊编译器前端最基础也最绕不开的一环——模型解析与导入。说白了，就是让编译器能“看懂”别人家的模型格式。

我刚开始做AI编译器那会儿，最头疼的就是模型格式不统一。ONNX、TFLite、Caffe……每个框架都有自己的“方言”。你想想看，一个模型从训练到部署，中间要经过多少道手？格式转换、算子映射、精度对齐，哪一步都可能出幺蛾子。

这一章，我就带你手撕ONNX和TFLite的解析流程，再聊聊自定义算子怎么注册进去。嗯，都是实战中踩过的坑。

4.1 ONNX模型格式解析

4.1.1 ONNX到底长什么样？

ONNX（Open Neural Network Exchange）本质上是一个protobuf序列化的文件。你用protoc反编译一下，就能看到它的内部结构。

核心数据结构就这几个：

ModelProto：最外层，包含版本号、生产者信息、计算图
GraphProto：计算图本体，包含节点、输入输出、初始化张量
NodeProto：每个算子节点，包含算子类型、输入输出名称、属性
TensorProto：权重、偏置等常量数据

我习惯把ONNX模型比作一个“菜谱”。GraphProto是整本菜谱，NodeProto是每个步骤，TensorProto就是食材本身。

关键点：ONNX的算子定义在onnx/onnx.proto中，所有标准算子都有明确的输入输出约束。但实际模型里经常出现“非标准”用法，这时候就需要你手动处理了。

4.1.2 手写一个ONNX解析器

我不会直接用现成的onnxruntime库，那样学不到东西。咱们从protobuf层面开始，自己写一个轻量解析器。

// 伪代码：ONNX模型解析核心流程
void parseONNXModel(const std::string& model_path) {
    // 1. 读取文件到内存
    std::ifstream file(model_path, std::ios::binary);
    std::string buffer((std::istreambuf_iterator<char>(file)), 
                        std::istreambuf_iterator<char>());
    
    // 2. 反序列化protobuf
    onnx::ModelProto model;
    model.ParseFromString(buffer);
    
    // 3. 获取计算图
    const auto& graph = model.graph();
    
    // 4. 遍历所有节点
    for (const auto& node : graph.node()) {
        // 算子类型
        std::string op_type = node.op_type();
        
        // 输入张量名称列表
        std::vector<std::string> inputs;
        for (const auto& name : node.input()) {
            inputs.push_back(name);
        }
        
        // 输出张量名称列表
        std::vector<std::string> outputs;
        for (const auto& name : node.output()) {
            outputs.push_back(name);
        }
        
        // 解析属性
        for (const auto& attr : node.attribute()) {
            if (attr.name() == "kernel_shape") {
                // 处理卷积核大小
            } else if (attr.name() == "strides") {
                // 处理步长
            }
            // ... 其他属性
        }
        
        // 5. 将节点注册到内部IR
        registerNode(op_type, inputs, outputs, attributes);
    }
    
    // 6. 加载权重
    for (const auto& tensor : graph.initializer()) {
        loadWeight(tensor.name(), tensor);
    }
}

这段代码看着简单，但实际坑不少。我曾经在一个项目中，遇到ONNX模型里某个节点的输入名称在initializer里找不到，排查了半天才发现是模型导出时有个权重被优化掉了。

小技巧：解析ONNX时，建议先打印graph.node()的数量和graph.initializer()的数量。如果节点数远大于初始化张量数，说明模型里有很多“非常量”输入，需要额外处理。

4.1.3 ONNX算子映射表

ONNX定义了上百个算子，但实际常用的就二三十个。我整理了一份常用映射表：

ONNX算子	功能描述	常见属性
Conv	卷积操作	kernel_shape, strides, pads, dilations
Relu	ReLU激活	无
Gemm	全连接层	alpha, beta, transA, transB
BatchNormalization	批归一化	epsilon, momentum
MaxPool	最大池化	kernel_shape, strides, pads
Reshape	张量变形	shape
Transpose	维度转置	perm

嗯，这里要注意：ONNX的Gemm算子和PyTorch的Linear层不是完全等价的。ONNX默认transB=1，意味着权重矩阵在乘法前需要转置。我第一次写映射时忘了这茬，结果推理结果全是错的。

4.2 TensorFlow Lite模型导入

4.2.1 TFLite的FlatBuffer结构

TFLite用的是Google的FlatBuffer序列化格式，和ONNX的protobuf完全不同。FlatBuffer的特点是“零反序列化成本”——你直接内存映射就能读取数据。

核心结构体：

Model：最外层，包含版本号和算子列表
SubGraph：子图，包含张量列表、算子列表、输入输出索引
Operator：算子节点，包含算子代码、输入输出张量索引
Tensor：张量定义，包含形状、类型、量化参数

我刚开始接触TFLite时，最不习惯的就是它用“索引”而不是“名称”来引用张量。ONNX里每个张量都有个字符串名字，TFLite里全是整数索引。调试起来特别费劲。

4.2.2 解析TFLite模型

// 伪代码：TFLite模型解析
void parseTFLiteModel(const std::string& model_path) {
    // 1. 内存映射文件
    int fd = open(model_path.c_str(), O_RDONLY);
    struct stat sb;
    fstat(fd, &sb);
    void* data = mmap(nullptr, sb.st_size, PROT_READ, MAP_PRIVATE, fd, 0);
    
    // 2. 获取FlatBuffer根对象
    const auto* model = tflite::GetModel(data);
    
    // 3. 获取子图（通常只有一个）
    const auto* subgraph = model->subgraphs()->Get(0);
    
    // 4. 遍历所有算子
    for (size_t i = 0; i < subgraph->operators()->size(); i++) {
        const auto* op = subgraph->operators()->Get(i);
        
        // 算子类型（内置算子代码）
        int op_code = model->operator_codes()->Get(op->opcode_index())->builtin_code();
        
        // 输入张量索引
        std::vector<int> inputs;
        for (size_t j = 0; j < op->inputs()->size(); j++) {
            inputs.push_back(op->inputs()->Get(j));
        }
        
        // 输出张量索引
        std::vector<int> outputs;
        for (size_t j = 0; j < op->outputs()->size(); j++) {
            outputs.push_back(op->outputs()->Get(j));
        }
        
        // 5. 根据算子代码创建内部IR节点
        createNodeFromOpCode(op_code, inputs, outputs);
    }
    
    // 6. 解析张量信息
    for (size_t i = 0; i < subgraph->tensors()->size(); i++) {
        const auto* tensor = subgraph->tensors()->Get(i);
        // 获取形状、类型、量化参数
    }
}

注意：TFLite的算子代码是枚举值，不同版本可能不一样。比如BuiltinOperator_CONV_2D在v1.x和v2.x中数值不同。解析时一定要检查模型版本号。

4.2.3 量化模型的特殊处理

TFLite最常用的场景是量化模型。量化参数存在每个张量的quantization字段里，包含scale和zero_point。

我遇到过最坑的情况：某个张量的量化参数是空的，但模型确实是量化模型。后来发现是TFLite转换时把某些“常量张量”的量化参数省略了，因为常量在编译期就已经量化好了。

处理量化模型时，我建议你这样做：

先检查tensor->quantization()->scale()是否为空
如果为空，检查该张量是否在initializer中
如果在，说明是常量，直接读取原始数据
如果不在，说明是动态量化，需要运行时计算

4.3 自定义算子注册

4.3.1 为什么需要自定义算子？

标准算子库再全，也总有覆盖不到的场景。比如你搞了个新的激活函数，或者硬件上有特殊指令加速某个操作。这时候就需要自定义算子。

我参与过一个项目，需要在NPU上实现一个“稀疏卷积”。ONNX和TFLite都没有这个算子，只能自己注册。

4.3.2 注册机制设计

一个通用的算子注册器，核心就三样东西：

算子名称：字符串标识，用于匹配
创建函数：根据输入参数创建算子实例
验证函数：检查输入输出是否合法

// 自定义算子注册器
class OpRegistry {
public:
    using CreatorFunc = std::function<std::shared_ptr<OpNode>(const OpDef&)>;
    using ValidatorFunc = std::function<bool(const OpDef&)>;
    
    // 注册算子
    void registerOp(const std::string& name, 
                    CreatorFunc creator, 
                    ValidatorFunc validator) {
        registry_[name] = {creator, validator};
    }
    
    // 创建算子
    std::shared_ptr<OpNode> createOp(const std::string& name, const OpDef& def) {
        auto it = registry_.find(name);
        if (it == registry_.end()) {
            throw std::runtime_error("Unknown op: " + name);
        }
        
        // 先验证
        if (!it->second.validator(def)) {
            throw std::runtime_error("Invalid op definition: " + name);
        }
        
        // 再创建
        return it->second.creator(def);
    }
    
private:
    std::unordered_map<std::string, std::pair<CreatorFunc, ValidatorFunc>> registry_;
};

// 使用示例：注册一个自定义算子
OpRegistry registry;
registry.registerOp("SparseConv", 
    [](const OpDef& def) -> std::shared_ptr<OpNode> {
        auto node = std::make_shared<SparseConvNode>();
        node->sparsity = def.getAttr<float>("sparsity");
        return node;
    },
    [](const OpDef& def) -> bool {
        return def.hasAttr("sparsity") && 
               def.getAttr<float>("sparsity") > 0.0f;
    }
);

设计建议：注册器最好做成单例模式，全局只有一个实例。这样在模型解析时，不管哪个模块都能访问到已注册的算子。

4.3.3 与模型解析的集成

自定义算子注册好后，怎么和ONNX/TFLite解析器联动呢？

我的做法是：解析器遇到未知算子时，不直接报错，而是先查注册器。如果注册器里有，就按自定义算子处理；如果没有，再报错。

// 在ONNX解析器中集成自定义算子
void parseONNXNode(const onnx::NodeProto& node, OpRegistry& registry) {
    std::string op_type = node.op_type();
    
    // 先查标准算子表
    if (isStandardOp(op_type)) {
        createStandardOp(op_type, node);
        return;
    }
    
    // 再查自定义算子注册器
    try {
        auto custom_op = registry.createOp(op_type, convertToOpDef(node));
        registerCustomOp(custom_op);
    } catch (const std::exception& e) {
        // 都找不到，报错
        throw std::runtime_error("Unsupported op: " + op_type + 
                                 ", " + e.what());
    }
}

这样做的好处是，你可以在不修改解析器核心代码的情况下，随时添加新的算子支持。我曾经在一个项目中，通过这种方式支持了十几种自定义算子，解析器一行代码都没改。

4.4 本章小结

模型解析是编译器前端的第一道坎。ONNX和TFLite虽然格式不同，但核心思路一致：把外部模型转换成编译器内部的统一IR。

我个人觉得，最难的不是解析本身，而是处理各种“边界情况”。比如ONNX里某个节点的输入名称拼写错误，或者TFLite量化参数缺失。这些坑，只有亲手踩过才能记住。

自定义算子注册机制，说白了就是给编译器开个“后门”。标准算子处理不了的，你自己写一个注册进去就行。但要注意，注册的算子一定要有完善的验证逻辑，否则后面优化阶段会出各种奇怪的问题。

好，这一章就到这里。下一章我们聊聊IR的中间表示设计，那是编译器的“灵魂”。

📚 课程目录

第四章：前端解析与模型导入

4.1 ONNX模型格式解析

4.1.1 ONNX到底长什么样？

4.1.2 手写一个ONNX解析器

4.1.3 ONNX算子映射表

4.2 TensorFlow Lite模型导入

4.2.1 TFLite的FlatBuffer结构

4.2.2 解析TFLite模型

4.2.3 量化模型的特殊处理

4.3 自定义算子注册

4.3.1 为什么需要自定义算子？

4.3.2 注册机制设计

4.3.3 与模型解析的集成

4.4 本章小结