深入理解MNN框架中的自定义算子开发

深入理解MNN框架中的自定义算子开发

MNN MNN is a blazing fast, lightweight deep learning framework, battle-tested by business-critical use cases in Alibaba 项目地址: https://gitcode.com/gh_mirrors/mn/MNN

前言

在深度学习框架中，算子(Operator)是最基础的计算单元，负责执行各种数学运算。MNN作为一款高性能的深度学习推理引擎，提供了完善的算子支持体系。本文将全面介绍在MNN框架中开发自定义算子的完整流程和技术要点。

一、MNN算子体系概述

MNN的算子体系采用分层设计，主要包含以下几个关键部分：

1. Schema层：定义算子的接口规范和数据格式
2. 转换层：负责将不同训练框架的算子转换为MNN内部表示
3. 计算层：实现算子的具体计算逻辑
4. 后端层：针对不同硬件平台的优化实现

这种分层设计使得MNN能够灵活支持多种训练框架和硬件平台，同时保持高效的推理性能。

二、自定义算子开发流程

2.1 前期准备

在开始开发自定义算子前，建议先检查MNN是否已经支持了该算子：

./MNNConvert -f CAFFE --OP
./MNNConvert -f TF --OP
./MNNConvert -f ONNX --OP 
./MNNConvert -f TORCH --OP

如果确认需要添加新算子，可以按照以下完整流程进行开发。

2.2 Schema描述定义

Schema是MNN算子的接口规范，需要在schema/default/MNN.fbs文件中定义：

1. 添加算子类型：在OpType枚举中追加新算子名称
2. 定义参数结构：如果算子需要参数，在OpParameter联合体和对应框架的fbs文件中添加参数描述

例如定义一个带参数的MyCustomOp：

enum OpType : int {
    // ...其他算子
    MyCustomOp
}

union OpParameter {
    // ...其他参数
    MyCustomOpParam
}

table MyCustomOpParam {
    padX:int;
    padY:int;
    kernelX:int;
    kernelY:int;
    strideX:int;
    strideY:int;
    dataType:DataType=DT_FLOAT;
}

修改完成后需要运行generate脚本重新生成头文件。

2.3 模型转换实现

根据原始模型格式，选择对应的转换器进行实现：

TensorFlow转换器示例

DECLARE_OP_CONVERTER(MyCustomOpTf);

void MyCustomOpTf::run(MNN::OpT* dstOp, TmpNode* srcNode, TmpGraph* tempGraph) {
    // 解析TF节点参数
    tensorflow::AttrValue value;
    find_attr_value(srcNode->tfNode, "pad_x", value);
    dstOp->main.AsMyCustomOpParam()->padX = value.i();
    // 其他参数解析...
}

MNN::OpType MyCustomOpTf::opType() {
    return OpType_MyCustomOp;
}

MNN::OpParameter MyCustomOpTf::type() {
    return OpParameter_MyCustomOpParam;
}

REGISTER_CONVERTER(MyCustomOpTf, MyCustomOp);

ONNX转换器示例

DECLARE_OP_CONVERTER(MyCustomOpOnnx);

void MyCustomOpOnnx::run(MNN::OpT* dstOp, 
                        const onnx::NodeProto* onnxNode,
                        std::vector<const onnx::TensorProto*> initializers) {
    // 解析ONNX节点属性
    for (const auto& attr : onnxNode->attribute()) {
        if (attr.name() == "pad_x") {
            dstOp->main.AsMyCustomOpParam()->padX = attr.i();
        }
        // 其他属性...
    }
}

2.4 维度计算实现

维度计算器负责推断输出Tensor的形状：

class MyCustomOpSizeComputer : public SizeComputer {
public:
    virtual bool onComputeSize(const MNN::Op* op, 
                             const std::vector<Tensor*>& inputs,
                             const std::vector<Tensor*>& outputs) const override {
        // 根据输入维度计算输出维度
        outputs[0]->buffer().dimensions = inputs[0]->dimensions();
        for (int i = 0; i < outputs[0]->dimensions(); ++i) {
            outputs[0]->setLength(i, inputs[0]->length(i));
        }
        return true;
    }
};

REGISTER_SHAPE(MyCustomOpSizeComputer, OpType_MyCustomOp);

2.5 计算实现

2.5.1 CPU实现

class CPUMyCustomOp : public Execution {
public:
    CPUMyCustomOp(Backend* backend) : Execution(backend) {}
    
    ErrorCode onResize(const std::vector<Tensor*>& inputs,
                     const std::vector<Tensor*>& outputs) override {
        // 准备计算所需的临时资源
        return NO_ERROR;
    }
    
    ErrorCode onExecute(const std::vector<Tensor*>& inputs,
                      const std::vector<Tensor*>& outputs) override {
        // 实现具体计算逻辑
        return NO_ERROR;
    }
};

REGISTER_CPU_OP_CREATOR(CPUMyCustomOpCreator, OpType_MyCustomOp);

2.5.2 GPU实现（以Metal为例）

class MetalMyCustomOp : public Execution {
public:
    MetalMyCustomOp(Backend* backend) : Execution(backend) {}
    
    ErrorCode onResize(const std::vector<Tensor*>& inputs,
                     const std::vector<Tensor*>& outputs) override {
        // 准备Metal计算资源
        return NO_ERROR;
    }
    
    void onEncode(const std::vector<Tensor*>& inputs,
                const std::vector<Tensor*>& outputs,
                id<MTLComputeCommandEncoder> encoder) override {
        // 编码Metal计算命令
    }
};

REGISTER_METAL_OP_CREATOR(MetalMyCustomOpCreator, OpType_MyCustomOp);

三、高级主题

3.1 几何计算优化

对于某些算子，可以通过实现几何计算来避免在各后端重复实现：

class GeometryMyCustomOp : public GeometryComputer {
public:
    virtual bool onCompute(const Op* op,
                         const std::vector<Tensor*>& inputs,
                         const std::vector<Tensor*>& outputs,
                         Context& context,
                         CommandBuffer& cmd) const override {
        // 使用现有算子组合实现新算子
        return true;
    }
};

REGISTER_GEOMETRY(GeometryMyCustomOp, OpType_MyCustomOp);

3.2 量化支持

对于需要支持量化的算子，需要额外实现量化版本：

class CPUMyCustomOpInt8 : public Execution {
    // 实现8位整型计算
};

REGISTER_CPU_OP_CREATOR(CPUMyCustomOpInt8Creator, OpType_MyCustomOpInt8);

四、最佳实践

1. 优先使用现有算子组合：能用现有算子组合实现的，就不要添加新算子
2. 全面测试：添加算子后要进行充分测试，包括不同输入形状、不同数据类型等
3. 性能优化：针对不同硬件平台进行针对性优化
4. 文档完善：为算子添加详细的文档说明

五、总结

MNN框架提供了完整的自定义算子开发体系，开发者可以根据需求灵活扩展算子支持。通过本文的介绍，相信读者已经掌握了在MNN中添加自定义算子的完整流程和关键技术要点。在实际开发中，建议参考MNN源码中现有算子的实现，遵循框架的设计规范，确保算子的性能和稳定性。

MNN MNN is a blazing fast, lightweight deep learning framework, battle-tested by business-critical use cases in Alibaba 项目地址: https://gitcode.com/gh_mirrors/mn/MNN