RKNN自定义算子开发详解：如何扩展NPU支持新操作

RKNN-Toolkit2是一款强大的神经网络模型转换和优化工具，它允许开发者将训练好的模型转换为适用于瑞芯微NPU的RKNN格式。然而，当遇到模型中包含NPU原生不支持的操作时，自定义算子功能就显得尤为重要。本文将详细介绍如何开发RKNN自定义算子，帮助开发者轻松扩展NPU的功能支持。## 什么是RKNN自定义算子？RKNN自定义算子是指用户根据特定需求实现的、NPU原生不支持的神经网络

苗素鹃Rich

994人浏览 · 2026-04-19 09:16:47

苗素鹃Rich · 2026-04-19 09:16:47 发布

RKNN自定义算子开发详解：如何扩展NPU支持新操作

【免费下载链接】rknn-toolkit2 项目地址: https://gitcode.com/gh_mirrors/rk/rknn-toolkit2

RKNN-Toolkit2是一款强大的神经网络模型转换和优化工具，它允许开发者将训练好的模型转换为适用于瑞芯微NPU的RKNN格式。然而，当遇到模型中包含NPU原生不支持的操作时，自定义算子功能就显得尤为重要。本文将详细介绍如何开发RKNN自定义算子，帮助开发者轻松扩展NPU的功能支持。

什么是RKNN自定义算子？

RKNN自定义算子是指用户根据特定需求实现的、NPU原生不支持的神经网络操作。通过开发自定义算子，开发者可以将各种先进的网络结构和操作部署到瑞芯微NPU上，极大地扩展了RKNN-Toolkit2的适用范围。

自定义算子的应用场景

自定义算子通常用于以下几种情况：

模型中包含NPU原生不支持的特殊操作
需要对某些操作进行特定优化以提高性能
实现论文中的最新算法和网络结构
满足特定业务场景的需求

RKNN自定义算子开发流程

开发RKNN自定义算子通常需要经过以下几个关键步骤：

1. 定义自定义算子类

首先需要定义一个继承自RKNN自定义算子基类的新类，并实现必要的方法。以下是一个简单的示例：

class cstDualResidual:
    # 自定义算子类型
    op_type = 'cstDualResidual'
    
    def shape_infer(self, node, in_shapes, in_dtypes):
        # 实现形状推断逻辑
        return in_shapes.copy(), in_dtypes.copy()
        
    def compute(self, node, inputs):
        # 实现算子计算逻辑
        x = inputs[0]
        y = inputs[1]
        alpha = get_node_attr(node, 'alpha')
        tmp_1 = x*alpha - y
        tmp_2 = y*alpha - x
        return [tmp_1, tmp_2]

2. 注册自定义算子

在加载模型之前，需要将自定义算子注册到RKNN对象中：

# 创建RKNN对象
rknn = RKNN(verbose=True)

# 注册自定义算子
ret = rknn.reg_custom_op(cstDualResidual())
if ret != 0:
    print('Register cstDualResidual op failed!')
    exit(ret)

3. 加载和转换模型

注册完成后，可以像处理普通模型一样加载和转换包含自定义算子的模型：

# 加载ONNX模型
ret = rknn.load_onnx(model=custom_model_path)
if ret != 0:
    print('Load model failed!')
    exit(ret)

# 构建模型
ret = rknn.build(do_quantization=False)
if ret != 0:
    print('Build model failed!')
    exit(ret)

# 导出RKNN模型
ret = rknn.export_rknn('dual_residual_custom.rknn')
if ret != 0:
    print('Export rknn model failed!')
    exit(ret)

4. 实现C++端推理代码

对于自定义GPU算子，还需要在C++端实现相应的推理代码。以下是一个简单的示例框架：

// 自定义算子实现
int32_t custom_argmax_gpu(rknn_context ctx, rknn_custom_op_t *op, const rknn_tensor_mem *inputs[], rknn_tensor_mem *outputs[]) {
    // 获取输入输出数据
    float *input = (float *)inputs[0]->virt_addr;
    float *output = (float *)outputs[0]->virt_addr;
    
    // 获取算子属性
    int axis = op->param[0].i;
    bool keepdims = op->param[1].b;
    
    // 实现GPU计算逻辑
    // ...
    
    return RKNN_SUCC;
}

// 注册自定义算子
rknn_custom_op_plugin custom_op_plugin[] = {
    {
        .op_name = "ArgMax",
        .op_type = RKNN_CUSTOM_OP_TYPE_GPU,
        .priority = 100,
        .compute = custom_argmax_gpu,
        .param_num = 2,
        .param = {
            {.name = "axis", .type = RKNN_CUSTOM_OP_PARAM_TYPE_INT32},
            {.name = "keepdims", .type = RKNN_CUSTOM_OP_PARAM_TYPE_BOOL},
        }
    },
    {NULL}
};

// 插件入口
RKNNCustomOpPluginInit(rknn_custom_op_plugin **plugin_list) {
    *plugin_list = custom_op_plugin;
    return 0;
}

5. 编译和部署插件

编译自定义算子插件，并将生成的库文件部署到目标设备：

# 编译插件
./build-linux.sh -t rk3588 -a aarch64 -b Release

# 推送插件到设备
adb push install/rknn_custom_gpu_op_demo_Linux/lib/librkcst_argmax.so /usr/lib/rknpu/op_plugins