RK3588部署RetinaFace:PyTorch到RKNN全流程与踩坑实录
1. 为什么要在 RK3588 上折腾 RetinaFaceRK3588 这颗芯片在边缘计算圈子里热度一直没降过。8 核 CPU4×A76 4×A55、Mali-G610 GPU、6 TOPS 算力的 NPU再加上丰富的外设接口拿来做视频结构化、人脸门禁、智能安防这类活儿非常合适。而 RetinaFace 作为人脸检测领域的经典模型精度高、能同时输出人脸框和五点关键点在 RK3588 上部署它基本是做人脸相关产品的入门必修课。但问题在于从 PyTorch 训练出来的模型到 RK3588 的 NPU 上跑起来中间隔着一条不短的链路PyTorch → ONNX → RKNN每一步都有坑。我自己前前后后折腾了好几块板子从 RK3568 到 RK3588踩过的坑能写满一页纸。这篇就把整个流程拆开揉碎讲清楚包括环境搭建、模型转换、量化校准、板端推理以及那些官方文档里不会写的细节。适合谁看如果你手上有一块 RK3588 开发板想跑人脸检测但卡在模型转换或者推理结果不对或者你正准备从零开始搭一套边缘人脸识别方案这篇内容应该能帮你省下不少时间。整个流程我会给出可直接复现的命令和代码参数选择也会解释清楚为什么这么定。2. 整体链路设计与方案选型2.1 从 PyTorch 到 RKNN 的完整链路先理清楚整条链路。RetinaFace 的原始实现是基于 PyTorch 的而 RK3588 的 NPU 只认 RKNN 格式的模型。所以中间必须经过 ONNX 这个中转站PyTorch 权重 → ONNX 模型 → RKNN 模型 → 板端推理每一步的核心任务不一样。PyTorch 到 ONNX 这一步重点是保证算子能被正确导出尤其是 RetinaFace 里用到的那些自定义操作比如 PriorBox 生成、NMS 等。ONNX 到 RKNN 这一步重点是算子兼容性检查和量化校准因为 NPU 对算子的支持是有限制的不是所有 ONNX 算子都能直接映射过去。为什么选 ONNX 做中转因为 RKNN Toolkit2 对 ONNX 的支持最成熟文档最全社区踩坑记录也最多。虽然理论上也支持 Caffe、TensorFlow但实际项目中 ONNX 是首选。另外ONNX 本身是个开放标准中间出了问题方便用 Netron 可视化排查这一点在调试阶段非常关键。2.2 为什么不用其他方案有人可能会问为什么不直接在板端用 PyTorch 推理答案很简单性能。RK3588 的 CPU 跑 PyTorch 推理RetinaFace 在 640×640 输入下大概只能到几帧而 NPU 跑 RKNN 模型可以轻松到 30 帧以上。功耗和发热也完全不是一个量级。那为什么不用其他推理框架比如 NCNN、MNN这些框架在 ARM CPU 上确实表现不错但它们用不到 RK3588 的 NPU。6 TOPS 的算力白白浪费掉做产品的时候成本就上去了。RKNN 是官方亲儿子对自家 NPU 的调度和优化是最到位的。还有一个选择是直接用 RKNN Model Zoo 里现成的 RetinaFace 模型。但实际项目中你往往需要自己训练或者微调模型比如换 backbone、改输入分辨率、调整 anchor 配置所以掌握完整的转换流程是必须的。2.3 环境版本选择的关键决策版本选择这块我要重点说一下因为这是最容易翻车的地方。RKNN Toolkit2 的版本和板端 NPU 驱动版本必须匹配否则会出现模型加载失败或者推理结果异常。我实测下来比较稳的组合是组件版本说明RKNN Toolkit21.6.0转换工具跑在 PC 上RKNPU2 Runtime1.6.0板端推理库NPU 驱动0.8.2 以上板端内核驱动Python3.8工具链对 3.8 支持最好ONNX1.12.0避免过高版本算子不兼容PyTorch1.13.0稳定版导出 ONNX 问题少注意RKNN Toolkit2 不建议用最新版新版本有时候会引入一些算子映射的变化导致原本能跑的模型突然报错。如果你不是非要新特性锁在 1.6.0 是比较稳妥的选择。PC 端建议用 Ubuntu 20.04 或 22.04不要用 Windows。虽然官方说支持 Windows但实际用下来各种路径问题和依赖冲突会让你怀疑人生。如果手头只有 Windows建议开个 WSL2但要注意 WSL2 里 USB 设备透传比较麻烦模型转换可以在 WSL2 里做板端调试还是得在原生 Linux 上。3. 核心细节解析与实操要点3.1 RetinaFace 模型结构的关键改动RetinaFace 原始实现里有一些操作是不适合直接导出 ONNX 的需要做改动。最常见的问题出在三个地方第一是 PriorBox 的生成。原始代码里 PriorBox 是在 forward 过程中动态计算的这会导致导出 ONNX 时出现动态 shape 的问题。解决办法是把 PriorBox 的计算提前到模型初始化阶段作为常量 buffer 存下来forward 里直接索引就行。第二是 NMS 操作。RetinaFace 推理后处理里有 NMS但 NMS 在 NPU 上支持不好。建议把 NMS 从模型里剥离出来放到 CPU 上做。模型只负责输出原始的框和分数后处理用 C 或者 Python 在 CPU 上实现。这样虽然多了一点 CPU 开销但换来了 NPU 的稳定运行。第三是输出层的处理。RetinaFace 有三个尺度的输出stride 8、16、32每个尺度输出 box 回归、分类分数、关键点回归。导出 ONNX 时要确保这些输出的 shape 是固定的不要有动态维度。具体改动示例基于常见的 RetinaFace 实现class RetinaFace(nn.Module): def __init__(self, cfgNone, phasetest): super(RetinaFace, self).__init__() self.phase phase self.cfg cfg # backbone、FPN、head 等定义省略 # 关键提前生成 priors 并注册为 buffer self.priors self._generate_priors() self.register_buffer(priors_buffer, self.priors) def _generate_priors(self): # 根据 feature map 尺寸和 anchor 配置生成 priors # 返回 shape 为 [num_priors, 4] 的 tensor pass def forward(self, x): # 正常前向传播 # 输出改为固定 shape不包含 NMS return loc, conf, landms提示改动完模型结构后一定要用相同的输入跑一遍 PyTorch 和 ONNX对比输出是否一致。差异超过 1e-4 就要检查是不是哪里改错了。3.2 ONNX 导出时的参数陷阱导出 ONNX 看起来就一行代码但里面的参数设置直接影响后续 RKNN 转换的成功率。torch.onnx.export( model, dummy_input, retinaface.onnx, opset_version11, input_names[input], output_names[loc, conf, landms], dynamic_axesNone, # 关键不要用动态轴 do_constant_foldingTrue, verboseFalse )几个关键点opset_version 选 11。这个版本对 RKNN 的兼容性最好。opset 12 以上有些算子 RKNN 不支持opset 9 以下又缺一些必要的算子。dynamic_axes 设为 None。RK3588 的 NPU 对动态 shape 支持有限固定 shape 能避免很多问题。如果你的应用需要多分辨率输入建议导出多个固定 shape 的模型而不是用一个动态模型。do_constant_folding 开启。这会把一些常量计算提前算好减小模型体积也能避免一些算子兼容问题。输入输出命名要清晰。后面 RKNN 转换和板端推理都要用到这些名字命名混乱容易搞错。导出完成后用 Netron 打开 ONNX 文件检查一下。重点看输入输出的 shape 对不对有没有奇怪的算子比如 Loop、If 这种控制流算子以及有没有 shape 为 dynamic 的维度。3.3 RKNN 转换配置的详细解读RKNN 转换是整个流程里最核心的一步。配置文件写得好不好直接决定模型能不能跑、跑得快不快。from rknn.api import RKNN rknn RKNN(verboseTrue) # 配置模型预处理 rknn.config( mean_values[[104, 117, 123]], std_values[[1, 1, 1]], target_platformrk3588, quantized_dtypeasymmetric_quantized-8, quantized_algorithmnormal, optimization_level3, compress_weightTrue, single_core_modeFalse ) # 加载 ONNX 模型 ret rknn.load_onnx(modelretinaface.onnx) if ret ! 0: print(Load ONNX failed) exit(ret) # 构建 RKNN 模型 ret rknn.build(do_quantizationTrue, dataset./dataset.txt) if ret ! 0: print(Build RKNN failed) exit(ret) # 导出 RKNN 模型 ret rknn.export_rknn(retinaface.rknn) if ret ! 0: print(Export RKNN failed) exit(ret)这里面的参数每一个都有讲究mean_values 和 std_values必须和训练时的预处理一致。RetinaFace 原始实现用的是 BGR 顺序mean 是 [104, 117, 123]。如果你训练时用的是 RGB 或者别的 mean这里要对应改。搞错了会导致检测结果完全不对但模型又能正常跑这种问题最难排查。quantized_dtype选asymmetric_quantized-8这是 RK3588 上精度和性能平衡最好的量化方式。如果对精度要求极高可以试试dynamic_fixed_point-16但模型体积会翻倍推理速度也会下降。quantized_algorithm选normal就行。mmse算法在某些模型上精度更好但转换时间会长很多而且不是所有模型都适用。optimization_level设 3这是最高优化级别。RKNN 会做一些算子融合和内存优化能提升推理速度。single_core_mode设 False让模型能用上 RK3588 的全部三个 NPU 核心。如果你的模型特别小用单核反而调度开销更小可以设 True 试试。3.4 量化校准数据集的准备量化校准是影响精度的关键环节。RKNN 的量化是 post-training quantizationPTQ需要一批代表性数据来统计激活值的分布。数据集准备有几个原则数量100 到 500 张就够了太多浪费时间太少统计不准。代表性要覆盖你实际应用场景的各种情况。做人脸检测的话要有不同光照、不同角度、不同大小的人脸最好还有没有人脸的背景图。格式图片要预处理成和模型输入一样的格式。如果模型输入是 640×640数据集里的图片也要 resize 到 640×640。路径文件dataset.txt 里每行一张图片的路径路径要写绝对路径避免相对路径找不到文件。# dataset.txt 示例 /home/user/calib/face_001.jpg /home/user/calib/face_002.jpg /home/user/calib/face_003.jpg ...注意校准图片不要用训练集里的图片否则量化后的精度评估会偏乐观。最好从验证集或者实际场景里采一批。4. 实操过程与核心环节实现4.1 PC 端环境搭建的完整步骤我习惯用 conda 管理环境避免和系统 Python 冲突。# 创建环境 conda create -n rknn python3.8 conda activate rknn # 安装 PyTorchCPU 版就够了转换不需要 GPU pip install torch1.13.0 torchvision0.14.0 --index-url https://download.pytorch.org/whl/cpu # 安装 ONNX pip install onnx1.12.0 onnxruntime1.14.0 # 安装 RKNN Toolkit2 pip install rknn-toolkit21.6.0 # 验证安装 python -c from rknn.api import RKNN; print(RKNN Toolkit2 installed)如果 pip 安装 RKNN Toolkit2 失败可以去官方仓库下载 whl 包手动安装。注意 whl 包要和 Python 版本、系统架构匹配。安装完成后建议跑一下官方提供的示例确认工具链能正常工作。官方示例通常在/usr/local/lib/python3.8/dist-packages/rknn/examples/下面。4.2 模型转换的完整脚本把前面的步骤串起来写一个完整的转换脚本import torch import numpy as np from rknn.api import RKNN # Step 1: 加载 PyTorch 模型 def load_pytorch_model(weight_path): from model import RetinaFace model RetinaFace() model.load_state_dict(torch.load(weight_path, map_locationcpu)) model.eval() return model # Step 2: 导出 ONNX def export_onnx(model, onnx_path, input_size640): dummy_input torch.randn(1, 3, input_size, input_size) torch.onnx.export( model, dummy_input, onnx_path, opset_version11, input_names[input], output_names[loc, conf, landms], dynamic_axesNone, do_constant_foldingTrue ) print(fONNX exported to {onnx_path}) # Step 3: 验证 ONNX 输出 def verify_onnx(model, onnx_path, input_size640): import onnxruntime as ort dummy_input torch.randn(1, 3, input_size, input_size) # PyTorch 输出 with torch.no_grad(): torch_out model(dummy_input) # ONNX 输出 sess ort.InferenceSession(onnx_path) onnx_out sess.run(None, {input: dummy_input.numpy()}) # 对比 for i, (t, o) in enumerate(zip(torch_out, onnx_out)): diff np.abs(t.numpy() - o).max() print(fOutput {i} max diff: {diff}) assert diff 1e-4, fOutput {i} diff too large # Step 4: 转换为 RKNN def convert_to_rknn(onnx_path, rknn_path, dataset_path): rknn RKNN(verboseTrue) rknn.config( mean_values[[104, 117, 123]], std_values[[1, 1, 1]], target_platformrk3588, quantized_dtypeasymmetric_quantized-8, quantized_algorithmnormal, optimization_level3, compress_weightTrue ) ret rknn.load_onnx(modelonnx_path) assert ret 0, Load ONNX failed ret rknn.build(do_quantizationTrue, datasetdataset_path) assert ret 0, Build RKNN failed ret rknn.export_rknn(rknn_path) assert ret 0, Export RKNN failed rknn.release() print(fRKNN exported to {rknn_path}) if __name__ __main__: model load_pytorch_model(retinaface.pth) export_onnx(model, retinaface.onnx) verify_onnx(model, retinaface.onnx) convert_to_rknn(retinaface.onnx, retinaface.rknn, ./dataset.txt)这个脚本跑完你会得到一个.rknn文件。但别急着往板子上扔先在 PC 上做一次仿真推理确认模型输出正常。4.3 PC 端仿真推理验证RKNN Toolkit2 提供了仿真推理功能可以在 PC 上模拟 NPU 的行为def simulate_inference(rknn_path, image_path): rknn RKNN() ret rknn.load_rknn(rknn_path) assert ret 0 ret rknn.init_runtime(targetrk3588) assert ret 0 # 加载并预处理图片 import cv2 img cv2.imread(image_path) img cv2.resize(img, (640, 640)) img img.astype(np.float32) # 推理 outputs rknn.inference(inputs[img]) for i, out in enumerate(outputs): print(fOutput {i} shape: {out.shape}, range: [{out.min():.4f}, {out.max():.4f}]) rknn.release() return outputs仿真推理的输出要和 ONNX 的输出对比。如果差异很大说明量化过程中精度损失太多需要调整量化参数或者增加校准数据。4.4 板端部署与推理板端需要安装 RKNPU2 Runtime。通常开发板的固件里已经带了如果没有可以从官方仓库编译安装。板端推理的 C 代码核心流程#include rknn_api.h // 加载模型 rknn_context ctx; int ret rknn_init(ctx, model_data, model_size, 0, nullptr); if (ret 0) { printf(rknn_init failed: %d\n, ret); return -1; } // 查询输入输出信息 rknn_input_output_num io_num; ret rknn_query(ctx, RKNN_QUERY_IN_OUT_NUM, io_num, sizeof(io_num)); // 设置输入 rknn_input inputs[1]; inputs[0].index 0; inputs[0].type RKNN_TENSOR_UINT8; inputs[0].size 640 * 640 * 3; inputs[0].fmt RKNN_TENSOR_NHWC; inputs[0].buf input_data; ret rknn_inputs_set(ctx, 1, inputs); // 推理 ret rknn_run(ctx, nullptr); // 获取输出 rknn_output outputs[3]; for (int i 0; i 3; i) { outputs[i].index i; outputs[i].want_float 1; } ret rknn_outputs_get(ctx, 3, outputs, nullptr); // 后处理NMS 等 // ... // 释放输出 rknn_outputs_release(ctx, 3, outputs); // 释放上下文 rknn_destroy(ctx);提示输入格式用RKNN_TENSOR_UINT8可以让 NPU 内部做归一化减少 CPU 预处理开销。但要注意 mean 和 std 的配置要和转换时一致。4.5 后处理 NMS 的实现要点前面说了 NMS 要放在 CPU 上做。RetinaFace 的后处理包括解码 box、过滤低分框、NMS、解码关键点。解码 box 的公式# priors: [num_priors, 4], 格式为 [cx, cy, w, h] # loc: [num_priors, 4], 网络输出的偏移量 # variances: [0.1, 0.2] boxes np.concatenate([ priors[:, :2] loc[:, :2] * variances[0] * priors[:, 2:], priors[:, 2:] * np.exp(loc[:, 2:] * variances[1]) ], axis1) # 转换为 [x1, y1, x2, y2] boxes[:, :2] - boxes[:, 2:] / 2 boxes[:, 2:] boxes[:, :2]NMS 用 OpenCV 的cv::dnn::NMSBoxes或者自己实现都可以。自己实现的话注意用 IoU 阈值 0.4 左右RetinaFace 原始实现用的是 0.4。关键点解码类似只是把 4 个偏移量换成 10 个5 个点 × 2 个坐标。5. 常见问题与排查技巧实录5.1 模型转换阶段的典型报错报错一E Catch exception when loading onnx model这个通常是因为 ONNX 模型里有 RKNN 不支持的算子。解决办法是用 Netron 打开 ONNX找到不支持的算子然后在 PyTorch 导出时替换掉。常见的坑包括Resize算子用opset 11的Resize而不是Upsample、HardSwishRKNN 1.6.0 支持不好换成ReLU或SiLU。报错二E Build model failed构建失败通常是量化校准的问题。检查 dataset.txt 里的路径是否正确图片是否能正常读取。另外如果校准图片的预处理和模型输入不匹配也会导致构建失败。报错三W The channel of input is not 3这个警告说明输入通道数不对。RetinaFace 输入是 3 通道检查 ONNX 的输入 shape 是不是[1, 3, 640, 640]。5.2 推理结果异常的排查思路现象一检测不到任何人脸先检查预处理。最常见的问题是 mean 和 std 搞反了或者 BGR/RGB 顺序错了。用一张确定有人脸的图片分别在 PyTorch、ONNX、RKNN 上跑对比输出。如果 PyTorch 和 ONNX 正常但 RKNN 不正常那就是量化的问题试试用dynamic_fixed_point-16或者增加校准数据。现象二框的位置偏移很大检查 priors 的生成逻辑。RKNN 转换后priors 是作为常量存在模型里的如果生成逻辑和训练时不一致框的位置就会偏。另外检查输入图片的 resize 方式是直接 resize 还是保持宽高比 padding这两种方式对应的后处理不一样。现象三关键点位置不对关键点的解码和 box 解码是独立的检查关键点的 variances 是不是设对了。RetinaFace 关键点的 variance 通常是[0.1, 0.2]和 box 一样。5.3 性能优化的实战经验经验一输入分辨率的选择640×640 是精度和速度的平衡点。如果对速度要求极高可以降到 320×320帧率能翻倍但小脸检测会变差。如果对精度要求高可以上到 1024×1024但帧率会掉到 10 帧以下。根据实际场景选。经验二多核 NPU 的利用RK3588 有三个 NPU 核心。如果模型比较大RKNN 会自动做多核调度。但如果模型很小单核跑反而更快因为多核调度的开销超过了并行收益。可以试试single_core_modeTrue和False的对比。经验三零拷贝推理板端推理时如果输入数据来自摄像头可以用 DMA 直接传到 NPU避免 CPU 拷贝。RK3588 的 MPP 和 RGA 模块可以配合 RKNN 做零拷贝但这部分需要改驱动层代码门槛较高。如果只是做原型验证先用普通方式跑通再说。5.4 常见问题速查表问题现象可能原因排查方法解决方案模型加载失败RKNN 版本和驱动不匹配查rknn_query返回的版本号统一版本到 1.6.0推理结果全为 0输入数据未正确设置打印输入数据范围检查rknn_inputs_set的 buf 和 size检测框偏移priors 不一致对比 PyTorch 和 RKNN 的 priors重新生成 priors 并注册为 buffer精度下降严重量化损失太大对比 ONNX 和 RKNN 输出增加校准数据或改用 16 位量化推理速度慢未用满 NPU 核心查 NPU 利用率调整single_core_mode和optimization_level内存泄漏输出未释放查内存增长确保rknn_outputs_release被调用注意板端调试时建议先用一张固定图片反复跑确认单帧结果正确后再接视频流。视频流引入的时序问题会让排查难度成倍增加。6. 从能跑到好用的进阶建议模型能在板子上跑起来只是第一步真正做产品还有很多细节要打磨。第一是模型剪枝和蒸馏。RetinaFace 原始模型对 RK3588 来说还是偏大如果追求极致性能可以用 MobileNet 或者更轻的 backbone 替换 ResNet50。精度会掉一些但速度能提升好几倍。实际项目中MobileNet 版本的 RetinaFace 在 RK3588 上跑 640×640 能到 60 帧以上。第二是多模型流水线。人脸检测只是第一步后面还有人脸对齐、特征提取、比对。RK3588 的 NPU 可以同时加载多个模型但要注意内存分配。建议把检测和对齐做成一个流水线检测到人脸后直接裁剪送对齐模型减少 CPU 和 NPU 之间的数据搬运。第三是温度控制。RK3588 满负荷跑 NPU 的时候发热不小如果做嵌入式产品散热设计要提前考虑。实测下来加个散热片能让 NPU 持续跑在最高频率不加的话跑几分钟就会降频。第四是固件和驱动的版本管理。量产的时候板端固件、NPU 驱动、RKNN Runtime 的版本要锁死不要随意升级。我见过太多因为升级了驱动导致模型跑不起来的案例。建议在项目初期就确定一套稳定版本然后冻结。第五是精度评估的标准化。不要只看几张图片的检测结果要跑完整的测试集算 mAP。量化后的模型 mAP 掉 1 到 2 个点是正常的如果掉超过 5 个点说明量化配置有问题需要重新调整。最后分享一个我踩过的坑RKNN 转换时如果开了compress_weightTrue模型体积会小很多但加载时间会变长。在需要快速启动的场景下可以关掉这个选项用空间换时间。这个参数在官方文档里只是一笔带过但实际影响不小建议根据你的启动时间要求来决定。