ARTICLE DETAIL

资讯详情

深耕网站建设与运营推广的一线实战洞察。

RKNN-Toolkit2 自定义 CPU 算子实战:以 cstSigmoid 替换 ONNX Sigmoid 的端到端流程(RKNPU2 运行时侧)

RKNN-Toolkit2 自定义 CPU 算子实战:以 cstSigmoid 替换 ONNX Sigmoid 的端到端流程(RKNPU2 运行时侧) 人工智能推理引擎模型量化模型优化边缘计算开发工具【免费下载链接】rknn-toolkit2项目地址https://gitcode.com/gh_mirrors/rk/rknn-toolkit2点击查看免费下载导读本文以 rknn-toolkit2 仓库中的rknn_custom_cpu_op_demo为完整范例系统讲解如何在 RKNPU2rknpu2运行时侧用 C 语言实现自定义 CPU 算子并以cstSigmoid 替换 ONNX 标准算子 Sigmoid为例串联起模型算子替换 → RKNN 模型转换 → 板端注册自定义算子 → AArch64 Linux / Android 交叉编译与部署的全链路。读完本文你将掌握rknn_register_custom_ops的注册机制、rknn_custom_op结构体的各回调字段语义、CPU 算子 compute 回调的写法以及 RK3562/RK356X/RK3576/RK3588/RV1126B 等平台的构建与运行方法。为什么需要自定义算子RKNN 编译器内置了大量算子支持但实际业务中常遇到两类情况一是 ONNX/PyTorch/TF 模型中的某个算子尚未被 RKNN 直接支持二是标准算子如 Sigmoid在特定量化或精度场景下不满足需求希望用自实现逻辑替代。RKNN 提供自定义算子Custom OP机制允许开发者绕过内置实现注册自己的算子逻辑。根据后端设备不同分为两类见 rknn_custom_op.hRKNN_TARGET_TYPE_CPU算子在 CPU 上执行本文主题RKNN_TARGET_TYPE_GPU算子以 OpenCL kernel 形式在 GPU 上执行对应仓库中的 rknn_custom_gpu_op_demo。需要特别说明的是该机制分为转换期与运行时两个阶段转换期由 rknn-toolkit2Python在将 ONNX 转 RKNN 时识别自定义算子并为其分配计算资源运行时则由 rknpu2 的 C/C API 把自定义算子的计算逻辑注册进rknn_context。本文的关联文档 rknn_custom_cpu_op_demo/README.md 聚焦运行时侧而转换期范例位于 rknn-toolkit2/examples/functions/custom_op/non-onnx_standard。端到端工作流程从模型替换到运行时注册自定义 CPU 算子的完整使用链路分四步修改 ONNX 图把原始 ONNX 模型中的Sigmoid节点类型改为自定义名cstSigmoid转换期注册在 rknn-toolkit2 中通过reg_custom_op注册同名自定义算子并导出 RKNN 模型运行时注册在板端 C 程序里rknn_init后调用rknn_register_custom_ops注册同名算子及计算回调推理RKNN 运行时遇到cstSigmoid节点时回调用户compute函数完成计算。第一步修改 ONNX 图把 Sigmoid 替换为 cstSigmoid转换期示例脚本 test.py 中的edit_onnx函数演示了替换逻辑遍历model.graph.node凡是op_type Sigmoid的节点一律改写为cstSigmoid再另存为yolox_s_custom.onnx。该范例以 YOLOX-S 检测模型为载体模型来源于 airockchip/YOLOX 开源项目仓库中已提供原始yolox_s.onnx供直接使用。def edit_onnx(in_model, output_model): import onnx model onnx.load(in_model) for node in model.graph.node: if node.op_type Sigmoid: node.op_type cstSigmoid onnx.save(model, output_model)第二步转换期注册自定义算子并导出 RKNN 模型继续沿用 test.py开发者需定义一个继承自 Python 的自定义算子类声明op_type、shape_infer与computeclass cstSigmoid: op_type cstSigmoid def shape_infer(self, node, in_shapes, in_dtypes): return in_shapes.copy(), in_dtypes.copy() def compute(self, node, inputs): return [1.0 / (1.0 np.exp(np.negative(inputs[0])))]随后在 RKNN 转换流程中依次调用rknn.config(...)、rknn.reg_custom_op(cstSigmoid())、rknn.load_onnx(modelcustom_model_path, ...)、rknn.build(do_quantizationTrue, datasetdataset.txt)与rknn.export_rknn(yolox_s_custom.rknn)。注意该脚本默认目标平台为rk3588实际部署时需按板卡平台修改rknn.config()的target_platform参数。转换完成后的 RKNN 模型即包含自定义算子cstSigmoid本仓库的 rknn_custom_cpu_op_demo/model 目录已按平台预置好转换产物yolox_s_custom_sigmoid.rknnRK3562 / RK3566_RK3568 / RK3576 / RK3588 / RV1126B。运行时自定义算子 API 与数据结构运行时侧的 API 全部定义在头文件 rknn_custom_op.h 中其依赖rknn_api.h使用前需同时#include rknn_api.h与#include rknn_custom_op.h。rknn_custom_op 核心结构注册的核心载体是rknn_custom_op结构体rknn_custom_op.h 中定义关键字段如下字段类型语义versionuint32_t自定义算子版本号targetrknn_target_type后端执行设备CPU 填RKNN_TARGET_TYPE_CPU值 1GPU 填RKNN_TARGET_TYPE_GPU值 2op_typechar[RKNN_MAX_NAME_LEN]算子类型名必须与模型中的节点类型一致如cstSigmoidinit回调函数指针[可选] 算子初始化回调在计算前被调用若返回RKNN_WARNING_SKIP_CUSTOM_OP_COMPUTE值为 -14且该算子类型被 RKNN 原生支持则回退到 RKNN 内置实现prepare回调函数指针[可选] 计算前的准备回调compute回调函数指针[必填] 算子核心计算回调自定义逻辑在此实现compute_native回调函数指针[可选] 原生属性计算回调当前版本不支持应置为nullptrdestroy回调函数指针[可选] 资源释放回调GPU 场景下还涉及cl_kernel_name、cl_kernel_source、cl_source_size、cl_build_options等 OpenCL 相关字段CPU 场景无需关心。相关辅助结构体rknn_custom_op_context自定义算子上下文包含target后端类型、internal_ctx算子内部上下文、gpu_ctxGPU OpenCL 上下文与priv_data用户私有数据指针rknn_custom_op_tensor算子输入/输出张量描述由attrrknn_tensor_attr描述形状、格式、量化参数与memrknn_tensor_mem描述虚拟地址、偏移等内存信息组成rknn_custom_op_attr算子属性查询结果name为属性名dtype为数组元素类型n_elems为元素个数data为属性数组指针。注册函数int rknn_register_custom_ops(rknn_context ctx, rknn_custom_op* op, uint32_t custom_op_num);按 rknn_custom_op.h 中的使用说明正确姿势是先创建并填充rknn_custom_op结构体数组把prepare/compute/compute_native/destroy回调挂到结构体上compute必填、其余可选、compute_native目前置nullptr在rknn_init之后调用rknn_register_custom_ops完成注册。另外还提供了rknn_custom_op_get_op_attr(op_ctx, attr_name, op_attr)用于在回调内按属性名获取算子属性以及get_custom_op_func类型的导出函数供.so动态库方式注册使用。示例源码剖析CPU 端自定义算子实现本节以 rknn_api_test_custom_cpu_op.cpp 为蓝本逐段分析。主程序执行流程main函数见 源码第 564 行起的命令行格式为Usage: rknn_custom_cpu_op_demo model_path [input_path] [loop_count]model_pathRKNN 模型路径必填input_path输入图片路径多个输入用#分隔loop_count推理循环次数用于性能测试默认 1。流程依次为读入 RKNN 模型 →rknn_init初始化上下文 →注册自定义算子→ 通过rknn_query查询 SDK/驱动版本、内存占用、输入输出属性与自定义字符串 → 加载并缩放输入图片load_image按输入张量要求的 H/W/C 用 stb_image 读图并stbir_resize_uint8缩放→rknn_inputs_set设置输入 → 循环rknn_run计时并输出每轮耗时与 FPS →rknn_outputs_get取输出 → 后处理与 NMS → 打印检测结果。注册代码注册块源码第 604-615 行展示了 CPU 算子的最小注册样板rknn_custom_op user_op[1]; memset(user_op, 0, sizeof(rknn_custom_op)); strncpy(user_op[0].op_type, cstSigmoid, RKNN_MAX_NAME_LEN - 1); user_op[0].version 1; user_op[0].target RKNN_TARGET_TYPE_CPU; user_op[0].compute compute_custom_sigmoid_float32; ret rknn_register_custom_ops(ctx, user_op, 1); if (ret 0) { printf(rknn_register_custom_op fail! ret %d\n, ret); return -1; }op_type与转换期注册的cstSigmoid保持一致target指明 CPU 后端compute指向自定义计算函数——三者构成注册的最小集合。compute 回调实现compute_custom_sigmoid_float32源码第 536-559 行是 float32 的 Sigmoid 实现核心要点是从inputs[0].mem.virt_addr inputs[0].mem.offset取得输入数据起始地址通过inputs[0].attr.dims与n_dims计算元素总数逐元素执行1 / (1 exp(-x))并写入输出张量int compute_custom_sigmoid_float32(rknn_custom_op_context* op_ctx, rknn_custom_op_tensor* inputs, uint32_t n_inputs, rknn_custom_op_tensor* outputs, uint32_t n_outputs) { unsigned char* in_ptr (unsigned char*)inputs[0].mem.virt_addr inputs[0].mem.offset; unsigned char* out_ptr (unsigned char*)outputs[0].mem.virt_addr outputs[0].mem.offset; const float* in_data (const float*)in_ptr; float* out_data (float*)out_ptr; int inside 1; for (int i 0; i inputs[0].attr.n_dims; i) { inside * inputs[0].attr.dims[i]; } for (int y 0; y inside; y) { out_data[y] 1 / (1 exp(-in_data[y])); } return 0; }由源码结构可推断自定义算子的内存访问完全建立在rknn_custom_op_tensor.mem.virt_addr/offset之上因此编写算子时务必严格按attr.dims计算张量尺寸并校验n_inputs/n_outputs同时保持 dtype 与模型量化配置一致本示例回调为 float32 形态。推理与后处理推理完成后示例复用 YOLOX 解码逻辑process_i8post_process对三个 stride 的输出特征图做反量化、阈值过滤BOX_THRESH 0.25、按类别 NMSNMS_THRESH 0.45最终打印class (left top right bottom) score格式的检测结果。类别名从./model/coco_80_labels_list.txt读取。AArch64 Linux 部署构建先设置交叉编译工具链前缀再调用构建脚本 build-linux.shexport GCC_COMPILER~/opt/gcc-linaro-7.5.0-2019.12-x86_64_aarch64-linux-gnu/bin/aarch64-linux-gnu ./build-linux.sh -t target -a arch -b build_type # 例如 ./build-linux.sh -t rk3588 -a aarch64 -b Release脚本参数说明来自脚本内帮助信息参数可选值说明-trk3562/rk3566/rk3568/rk3576/rk3588/rv1126b目标 SoCrk356x等价于RK3566_RK3568-aaarch64/armhf目标架构-bDebug/Release构建类型默认Release脚本内部会将小写 target 映射为模型目录名如rk3588→RK3588并基于GCC_COMPILER生成CC/CXX通过 CMake 交叉编译后make install。CMake 逻辑见 CMakeLists.txt其从../../../runtime/Linux/librknn_api引入librknnrt.so与rknn_api.h/rknn_custom_op.h头文件安装时把可执行文件、librknnrt.so、模型图片、标签文件及对应平台的model/TARGET_SOC目录一并装入install/rknn_custom_cpu_op_demo_Linux。安装把install/rknn_custom_cpu_op_demo_Linux拷贝到设备的/userdata/下。若使用瑞芯微 EVB 板可通过 adb 推送adb push install/rknn_custom_cpu_op_demo_Linux/ /data/若板卡启用了 sshd 服务也可用 scp 等方式传输程序与 RKNN 模型。运行adb shell cd /data/rknn_custom_cpu_op_demo_Linux export LD_LIBRARY_PATH./lib ./rknn_custom_cpu_op_demo model/TARGET_PLATFORM/yolox_s_custom_sigmoid.rknn model/test_image.jpg 1其中TARGET_PLATFORM对应模型目录名RK3562/RK3566_RK3568/RK3576/RK3588/RV1126B最后的1是推理循环次数。由于librknnrt.so安装在相对lib目录必须先导出LD_LIBRARY_PATH。Android 部署构建先导出ANDROID_NDK_PATH再调用 build-android.shexport ANDROID_NDK_PATH~/opts/ndk/android-ndk-r18b ./build-android.sh -t target -a arch [-b build_type] # 例如 ./build-android.sh -t rk3568 -a arm64-v8a -b Release脚本参数与 Linux 版基本一致差异点在于参数可选值说明-trk3562/rk3566/rk3568/rk3576/rk3588Android 侧暂不支持rv1126b-aarm64-v8a/armeabi-v7aAndroid ABI-bDebug/Release构建类型默认Release脚本要求 NDK r18/r19 等版本脚本注释明确NDK Version r18, r19 is recommended内部以clang工具链、-DANDROID_STLc_static、-DANDROID_PLATFORMandroid-24调用 CMake产物安装到install/rknn_custom_cpu_op_demo_Android。安装与运行adb push install/rknn_custom_cpu_op_demo_Android/ /data/ adb shell cd /data/rknn_custom_cpu_op_demo_Android/ export LD_LIBRARY_PATH./lib ./rknn_custom_cpu_op_demo model/TARGET_PLATFORM/yolox_s_custom_sigmoid.rknn model/test_image.jpg 1运行结果解读RK3588以 RK3588 平台为例对model/test_image.jpg的检测输出如下bus (92 132 558 438) 0.933 person (104 237 221 534) 0.901 person (211 240 285 504) 0.901 person (472 248 559 525) 0.774 person (79 329 118 512) 0.347 bicycle (80 420 96 515) 0.253可见使用自定义cstSigmoid算子后YOLOX 模型仍能正确输出 bus、person、bicycle 等目标的边框与置信度说明自定义 CPU 算子成功参与了整图推理且不影响后处理链路。注意事项与常见问题算子名一致性转换期reg_custom_op注册的op_typePython 类cstSigmoid.op_type必须与运行时rknn_register_custom_ops中user_op[0].op_type完全一致否则模型运行时找不到对应算子。平台差异不同平台、不同版本的工具链与驱动可能产生略有差异的结果README 已注明该前提。回退机制init回调若返回RKNN_WARNING_SKIP_CUSTOM_OP_COMPUTE-14且该算子类型恰好被 RKNN 原生支持运行时将回退使用内置实现可用于排查自定义实现精度问题。CPU/GPU 后端选择RKNN_TARGET_TYPE_CPU走本文的纯 C 回调若算子需要 OpenCL 加速应参考 rknn_custom_gpu_op_demo 并配置cl_kernel_*字段。量化形态本文示例的 compute 回调为 float32 实现若模型为 INT8 量化需要在回调内自行做反量化/再量化可参考示例中qnt_f32_to_affine与deqnt_affine_to_f32辅助函数或借助rknn_custom_op_tensor.attr中的zp/scale字段处理。环境变量Linux 与 Android 运行前都必须export LD_LIBRARY_PATH./lib否则动态链接器找不到板端librknnrt.so。小结通过rknn_custom_cpu_op_demo这份开箱即用的范例开发者可以完整掌握 RKNN 自定义 CPU 算子的全生命周期用 test.py 在转换期替换算子并导出 RKNN再用 rknn_api_test_custom_cpu_op.cpp 在运行时注册cstSigmoid的 float32 计算逻辑最终通过 build-linux.sh / build-android.sh 交叉编译并部署到 RK3562、RK3566_RK3568、RK3576、RK3588、RV1126B 等目标平台。这套改图 注册 回调的模式可推广到任意需要定制算子逻辑的场景是 RKNPU2 生态中处理非标准算子问题的标准解法。赞分享人工智能推理引擎模型量化模型优化边缘计算开发工具【免费下载链接】rknn-toolkit2项目地址https://gitcode.com/gh_mirrors/rk/rknn-toolkit2点击查看免费下载相关推荐Wox Shell 插件完全指南在启动器中直接执行命令、管理历史与工作目录Wox Shell 插件完全指南在启动器中直接执行命令、管理历史与工作目录 本指南围绕 Wox 内置系统插件 Shell 展开讲解如何用 触发关键字在启动器人工智能推理引擎模型量化模型优化边缘计算开发工具rknn-toolkit2 实战将 YOLOv5 ONNX 模型转换为 RKNN 模型并部署到 RKNPUrknn toolkit2 实战将 YOLOv5 ONNX 模型转换为 RKNN 模型并部署到 RKNPU 本篇技术指南围绕 RKNN Toolkit2 在人工智能推理引擎模型量化模型优化边缘计算开发工具RKNN-Toolkit2 开发指南读懂 Rockchip RKNN 软件栈从模型转换到 NPU 端侧部署RKNN Toolkit2 开发指南读懂 Rockchip RKNN 软件栈从模型转换到 NPU 端侧部署 RKNNRockchip Neural Net人工智能推理引擎模型量化模型优化边缘计算开发工具上一篇CopilotKit AG2 集成Pre-Built Sidebar 预构建侧边栏的端到端验证指南下一篇oh-my-pi 内置规则解析如何在测试中彻底告别真实定时器让 CI 不再 flaky创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考
返回列表