ARTICLE DETAIL

资讯详情

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

RetinaFace在RK3588上的部署实战:从PyTorch到RKNN的完整指南

RetinaFace在RK3588上的部署实战:从PyTorch到RKNN的完整指南 1. 为什么RetinaFace在RK3588上值得折腾RetinaFace 这个模型在人脸检测领域算是老牌劲旅了。它基于 RetinaNet 架构在特征金字塔的每一层上都做密集锚框预测同时输出人脸框、五点关键点双眼、鼻尖、左右嘴角以及置信度。相比 MTCNN 那种级联结构RetinaFace 一步到位精度高、速度快在 WIDER FACE 榜单上长期霸榜。而 RK3588 这颗芯片八核 CPU4×A76 4×A55加上 6 TOPS 算力的 NPU跑 RetinaFace 这种量级的模型完全绰绰有余。但问题在于PyTorch 训练出来的模型不能直接扔到 RK3588 的 NPU 上跑。RK3588 的 NPU 只认 RKNN 格式中间需要经过 ONNX 导出、RKNN 转换、量化、板端推理这一整套流程。每一步都有坑而且坑和坑之间还互相牵连。我前后在 RK3588 上部署过好几个视觉模型RetinaFace 算是其中比较典型的一个——它既有 backbone 的卷积网络又有 FPN 的多尺度特征融合还有 SSH 上下文模块和最后的检测头结构不算简单转换过程中容易出问题的地方不少。这篇内容适合两类人一类是手里已经有 RK3588 开发板、想把 PyTorch 模型搬到板子上跑的嵌入式 AI 工程师另一类是刚开始接触 RKNN 工具链、想找一个完整案例练手的开发者。我会从环境搭建开始一步步走到板端推理把每个环节的坑都标出来包括我实际踩过的和社区里反馈比较多的。注意整个流程涉及 PC 端x86 Ubuntu和板端RK3588两个环境PC 端负责模型转换板端负责推理。两边的 Python 环境和依赖版本要分开管理不要混在一起。2. 环境搭建PC端和板端的依赖怎么配才不打架2.1 PC端RKNN-Toolkit2的安装与版本匹配RKNN-Toolkit2 是瑞芯微官方提供的模型转换工具运行在 x86 Linux 上。它的版本和板端 RKNPU2 驱动版本必须匹配否则转换出来的 RKNN 模型在板子上加载会直接报错。我用的组合是 RKNN-Toolkit2 1.6.0 RKNPU2 1.6.0这个组合在 RK3588 上验证过比较稳定。安装方式推荐用 conda 建一个独立环境Python 版本选 3.8 或 3.10 都行但要注意 RKNN-Toolkit2 1.6.0 对 Python 3.10 的支持在部分发行版上会有 numpy 版本冲突。我个人的做法是conda create -n rknn python3.8 conda activate rknn pip install rknn-toolkit21.6.0 pip install onnx1.14.0 pip install torch1.13.1 torchvision0.14.1这里有个细节onnx 的版本不要装太新。我有一次手贱装了 onnx 1.16结果导出 RetinaFace 的时候报了一堆算子不支持的错误。退回到 1.14 就正常了。PyTorch 用 1.13.1 是因为 RetinaFace 的官方实现biubug6/Pytorch_Retinaface在这个版本上验证过换成 2.x 可能会有 API 变动导致加载权重失败。另外RKNN-Toolkit2 在转换时内部会调用 onnx 的优化器如果你的 onnx 版本和 toolkit 内置的期望版本不一致可能会在rknn.load_onnx()这一步就卡住。所以装完之后最好跑一个简单的测试from rknn.api import RKNN rknn RKNN() print(rknn.version())能打印出版本号就说明基础环境没问题。2.2 板端RKNPU2运行环境的部署板端需要安装 RKNPU2 的运行时库和 Python 接口。如果你烧录的是瑞芯微官方的 Ubuntu 20.04 固件系统里通常已经带了 NPU 驱动但 Python 的 rknn-toolkit-lite2 需要自己装。# 板端执行 pip install rknn-toolkit-lite21.6.0注意板端的 Python 版本要和 PC 端转换时用的版本尽量一致虽然 RKNN 模型本身是跨版本的但 rknn-toolkit-lite2 的 wheel 包对 Python 版本有要求。RK3588 官方固件一般是 Python 3.8所以 PC 端也用 3.8 最省事。验证板端 NPU 是否可用cat /sys/kernel/debug/rknpu/version如果输出类似RKNPU2: 0.9.6之类的版本信息说明驱动正常。如果这个文件不存在可能需要更新内核或检查设备树里 NPU 节点是否使能。提示板端推理时如果报librknnrt.so: cannot open shared object file说明 RKNPU2 运行时库没装或者路径不对。可以手动把 librknnrt.so 拷贝到/usr/lib/下然后ldconfig刷新一下。2.3 两个环境之间的文件传输PC 端转换出来的 .rknn 文件需要传到板子上。最方便的方式是用 scp 或者 adb push。如果板子通过 USB 连接adb 也可以用adb push retinaface.rknn /userdata/ adb push test.jpg /userdata/如果板子已经联网scp 更直接scp retinaface.rknn root192.168.1.100:/userdata/我一般会在板子上建一个/userdata/retinaface/目录把模型、测试图片、推理脚本都放进去方便管理。3. 从PyTorch到ONNX导出环节的暗礁3.1 RetinaFace模型结构的特殊性RetinaFace 的官方 PyTorch 实现以 biubug6 的版本为例在推理时会做几件事输入一张 RGB 图片经过 backbone通常是 ResNet50 或 MobileNet0.25提取特征然后通过 FPN 做多尺度融合再经过 SSH 模块增强感受野最后在三个尺度上分别输出分类置信度和框回归结果。这个结构在导出 ONNX 时有两个麻烦点。第一模型 forward 函数里可能有 Python 层面的控制流比如根据输入尺寸动态调整 anchor 生成。ONNX 导出时如果遇到这种动态逻辑要么固定输入尺寸要么改写 forward 函数。第二RetinaFace 的输出不是单个 tensor而是多个 tensor 的列表每个尺度一组导出后 ONNX 的输出节点会有多个RKNN 转换时要确保所有输出都被正确识别。我的做法是在导出之前先把模型包装一层固定输入尺寸为 640×640并且把 forward 的返回值整理成一个 tuple方便后续处理。import torch import torch.nn as nn from models.retinaface import RetinaFace class RetinaFaceWrapper(nn.Module): def __init__(self, model): super().__init__() self.model model def forward(self, x): # 原始模型返回 [loc, conf, landmarks] out self.model(x) return out[0], out[1], out[2]这里out[0]是框回归out[1]是置信度out[2]是关键点。把它们分开返回ONNX 导出后就是三个独立的输出节点RKNN 转换时更容易对应。3.2 ONNX导出的参数设置与验证导出命令本身不复杂但参数要设对model RetinaFace(cfgcfg, phasetest) model.load_state_dict(torch.load(RetinaFace-ResNet50.pth, map_locationcpu)) model.eval() wrapper RetinaFaceWrapper(model) dummy_input torch.randn(1, 3, 640, 640) torch.onnx.export( wrapper, dummy_input, retinaface.onnx, opset_version11, input_names[input], output_names[loc, conf, landmarks], dynamic_axesNone # 固定尺寸不用动态轴 )opset_version 选 11 是因为 RKNN-Toolkit2 对 opset 11 的支持最完善。选 12 或 13 可能会遇到某些算子不被识别的问题。dynamic_axes 设为 None 是因为 RK3588 的 NPU 对动态尺寸支持有限固定 640×640 最稳妥。导出之后一定要用 onnxruntime 验证一下import onnxruntime as ort import numpy as np sess ort.InferenceSession(retinaface.onnx) input_name sess.get_inputs()[0].name dummy np.random.randn(1, 3, 640, 640).astype(np.float32) outputs sess.run(None, {input_name: dummy}) print([o.shape for o in outputs])如果输出 shape 符合预期比如 loc 是 [1, 16800, 4]conf 是 [1, 16800, 2]landmarks 是 [1, 16800, 10]说明导出成功。如果 shape 不对或者报错大概率是 forward 函数里有动态逻辑没处理好。注意RetinaFace 的 anchor 数量取决于输入尺寸和 FPN 层数。640×640 输入下三个尺度的 anchor 总数是 16800。如果你换了输入尺寸这个数字会变RKNN 转换后的后处理代码也要跟着改。3.3 常见导出错误与修复我遇到过几种典型的导出错误这里列一下错误信息原因修复方法RuntimeError: ONNX export failed: Couldnt export operator aten::xxx模型里有 ONNX 不支持的 PyTorch 算子改写 forward用 ONNX 支持的算子替代TypeError: forward() missing 1 required positional argument模型 forward 需要额外参数用 wrapper 包装固定额外参数输出 shape 为 [0] 或 None动态控制流导致固定输入尺寸去掉 if/else 分支TracerWarning: Converting a tensor to a Python boolean模型里有 tensor 比较操作改写为 torch.where 或固定分支最常见的是最后一种。RetinaFace 的原始代码里可能有类似if x.shape[2] 100:的判断这种在 trace 模式下会出问题。解决办法是把这些判断去掉或者用torch.jit.script代替 trace。但 RKNN 转换更推荐 trace 模式所以还是改代码更直接。4. RKNN转换量化、算子支持和精度调优4.1 RKNN模型转换的完整脚本ONNX 准备好之后就可以用 RKNN-Toolkit2 转成 RKNN 了。完整脚本如下from rknn.api import RKNN rknn RKNN(verboseTrue) # 配置 rknn.config( mean_values[[123.675, 116.28, 103.53]], std_values[[58.395, 57.12, 57.375]], target_platformrk3588, quantized_dtypeasymmetric_quantized-8, optimization_level3 ) # 加载ONNX ret rknn.load_onnx(modelretinaface.onnx) if ret ! 0: print(Load ONNX failed) exit(ret) # 构建 ret rknn.build(do_quantizationTrue, datasetdataset.txt) if ret ! 0: print(Build failed) exit(ret) # 导出 ret rknn.export_rknn(retinaface.rknn) if ret ! 0: print(Export failed) exit(ret) rknn.release()这里的 mean_values 和 std_values 要和 PyTorch 训练时的预处理一致。RetinaFace 官方用的是 ImageNet 的均值和方差但注意顺序是 RGB而 RKNN 默认按 RGB 处理所以直接填就行。如果你的模型训练时用的是 BGR那这里要调整顺序。4.2 量化数据集准备为什么100张图比1000张更管用量化是 RKNN 转换里最影响精度的环节。RKNN-Toolkit2 用的是 PTQ训练后量化需要一个校准数据集来统计激活值的分布。数据集不用多100 到 200 张就够但一定要有代表性。我一开始偷懒用了 20 张同一场景的图片结果量化后模型在其他人脸图片上漏检严重。后来换成 200 张涵盖不同光照、不同角度、不同人脸的图片精度就回来了。数据集的组织方式很简单建一个dataset.txt每行一张图片的路径./calib/face_001.jpg ./calib/face_002.jpg ... ./calib/face_200.jpg图片尺寸不用和模型输入一致RKNN 会自动 resize。但建议校准图片的尺寸分布和实际推理场景接近比如都是 640×480 左右的监控画面截图。提示校准图片里最好包含一些人脸目标不要全是纯背景。如果校准集里没有人脸量化后的模型对人脸区域的激活值统计不准检测精度会下降。4.3 算子支持与自定义层处理RetinaFace 用到的算子大部分 RKNN 都支持但有几个需要注意Resize/InterpolateFPN 上采样用的RKNN 支持 nearest 和 bilinear但 bilinear 在某些版本上需要指定half_pixel_centers参数。ConcatFPN 融合用的支持没问题。Softmax分类头用的支持。PriorBox如果模型里包含了 PriorBox 层RKNN 可能不支持需要把 anchor 生成放到后处理里做。我建议在导出 ONNX 时就把 PriorBox 去掉只保留纯网络输出。anchor 生成和解码放在板端的 Python 后处理里做这样更灵活也避免 RKNN 转换时算子不支持的问题。如果转换时遇到Unsupported operator错误可以用rknn.config(optimization_level0)先关掉优化看看是哪个算子的问题。然后要么改写模型要么用 RKNN 的自定义算子接口但 RK3588 对自定义算子的支持有限能改写就改写。4.4 量化精度验证与调优转换完成后RKNN-Toolkit2 提供了精度分析工具rknn.accuracy_analysis( inputs[test.jpg], output_dir./snapshot, targetrk3588 )这个命令会生成每一层的量化误差报告放在snapshot目录下。重点看error_analysis.txt里面会列出哪些层的余弦相似度低于阈值一般 0.99 以上算正常。如果某些层误差特别大可以尝试把这些层加入hybrid_quantization列表让它们保持浮点计算。增加校准图片数量。调整quantized_dtype比如从asymmetric_quantized-8换成dynamic_fixed_point-8。不过混合量化会降低推理速度因为浮点层要在 CPU 上跑。我的经验是RetinaFace 的 backbone 部分量化误差通常很小FPN 和检测头部分偶尔会有问题。如果精度下降在可接受范围内比如 mAP 掉 1-2 个点就不用折腾混合量化了。5. 板端推理后处理才是真正的战场5.1 RKNN模型加载与推理接口板端的推理脚本用 rknn-toolkit-lite2 写from rknnlite.api import RKNNLite import cv2 import numpy as np rknn_lite RKNNLite() ret rknn_lite.load_rknn(retinaface.rknn) ret rknn_lite.init_runtime(core_maskRKNNLite.NPU_CORE_0) img cv2.imread(test.jpg) img cv2.cvtColor(img, cv2.COLOR_BGR2RGB) img cv2.resize(img, (640, 640)) img img.astype(np.float32) img (img - [123.675, 116.28, 103.53]) / [58.395, 57.12, 57.375] img np.expand_dims(img, axis0) outputs rknn_lite.inference(inputs[img])core_mask可以指定用哪个 NPU 核心。RK3588 有三个 NPU 核心可以并行跑三个模型实例。如果只跑一个模型用NPU_CORE_0就行。如果要跑多路视频可以用NPU_CORE_0_1_2让三个核心一起上。5.2 Anchor生成与解码的板端实现前面说了PriorBox 被去掉了所以 anchor 要在板端生成。RetinaFace 的 anchor 生成逻辑是对于每个特征层每个像素点生成 2 个 anchor一个正方形一个长方形然后乘以步长和缩放系数。def generate_anchors(height, width, step, scales): anchors [] for y in range(height): for x in range(width): cx (x 0.5) * step cy (y 0.5) * step for scale in scales: anchors.append([cx, cy, scale, scale]) return np.array(anchors)640×640 输入下三个特征层的尺寸分别是 80×80、40×40、20×20对应的步长是 8、16、32。每个位置 2 个 anchor总数是 (80×80 40×40 20×20) × 2 16800。这个数字要和 ONNX 输出的 shape 对上。解码就是把网络输出的偏移量应用到 anchor 上def decode(loc, anchors, variances[0.1, 0.2]): boxes np.concatenate([ anchors[:, :2] loc[:, :2] * variances[0] * anchors[:, 2:], anchors[:, 2:] * np.exp(loc[:, 2:] * variances[1]) ], axis1) boxes[:, :2] - boxes[:, 2:] / 2 boxes[:, 2:] boxes[:, :2] return boxes关键点解码类似只是每个关键点有两个坐标x, y需要分别解码。5.3 NMS与关键点后处理的性能优化NMS非极大值抑制是后处理里最耗时的部分。16800 个框如果直接用 Python 循环做 NMS在 RK3588 的 CPU 上可能要跑几十毫秒。优化方法有几个先做置信度过滤把置信度低于阈值比如 0.5的框先扔掉通常能过滤掉 90% 以上的框。用 NumPy 向量化实现 NMS避免 Python 循环。限制每类最多保留的框数比如最多 750 个超过的直接按置信度排序截断。def nms(dets, thresh): x1 dets[:, 0] y1 dets[:, 1] x2 dets[:, 2] y2 dets[:, 3] scores dets[:, 4] areas (x2 - x1 1) * (y2 - y1 1) order scores.argsort()[::-1] keep [] while order.size 0: i order[0] keep.append(i) xx1 np.maximum(x1[i], x1[order[1:]]) yy1 np.maximum(y1[i], y1[order[1:]]) xx2 np.minimum(x2[i], x2[order[1:]]) yy2 np.minimum(y2[i], y2[order[1:]]) w np.maximum(0.0, xx2 - xx1 1) h np.maximum(0.0, yy2 - yy1 1) inter w * h ovr inter / (areas[i] areas[order[1:]] - inter) inds np.where(ovr thresh)[0] order order[inds 1] return keep这个 NumPy 版本的 NMS 在 RK3588 上跑 16800 个框大概 5-10 毫秒可以接受。如果还嫌慢可以考虑用 C 实现或者用 OpenCV 的cv2.dnn.NMSBoxes。关键点后处理就是把解码后的关键点坐标映射回原图尺寸。注意 RKNN 推理时图片被 resize 到了 640×640后处理时要按比例还原。5.4 实测性能数据与调优建议我在 RK3588 上实测的数据如下环节耗时ms图片预处理resize 归一化3-5NPU 推理RetinaFace ResNet5025-35Anchor 生成 解码2-3NMS5-10关键点后处理1-2总计36-55这个速度跑单路视频30 FPS完全够用。如果换成 MobileNet0.25 的 backboneNPU 推理可以降到 8-12 毫秒整体 20 毫秒以内。调优建议如果多路视频用NPU_CORE_0_1_2让三个核心并行但要注意内存带宽瓶颈。预处理可以用 RGARockchip Graphics Accelerator硬件加速比 OpenCV 的 resize 快很多。如果对精度要求不高可以把输入尺寸从 640×640 降到 320×320推理速度翻倍但小脸检测会变差。注意RK3588 的 NPU 在跑量化模型时如果输入数据分布和校准集差异很大可能会出现精度骤降。建议在实际部署前用真实场景的视频流做一轮测试看看有没有漏检或误检。6. 那些让我熬夜的坑和最终解决方案6.1 模型加载失败版本不匹配的连锁反应我最开始用的是 RKNN-Toolkit2 1.5.0 转的模型板端装的是 RKNPU2 1.6.0结果加载时报RKNN model version mismatch。后来统一成 1.6.0 就好了。这个问题的坑在于PC 端和板端的版本号不一定要求完全一致但主版本号必须相同。比如 1.6.0 转的模型板端 1.6.2 可以加载但 1.5.x 就不行。另一个相关问题是 Python 版本。PC 端用 Python 3.10 转的模型板端 Python 3.8 加载时可能会报pickle相关的错误。虽然 RKNN 模型本身是二进制格式但 rknn-toolkit-lite2 在加载时会用 pickle 反序列化一些元数据Python 版本差异可能导致兼容性问题。所以两边都用 3.8 最保险。6.2 量化后精度暴跌校准集的锅前面提过我一开始用 20 张同场景图片做校准量化后模型在其他人脸图片上几乎全漏。后来分析accuracy_analysis的报告发现 FPN 的某些 Concat 层余弦相似度只有 0.85。换成 200 张多样化图片后相似度回升到 0.98 以上检测精度也恢复正常。这个坑的教训是校准集的质量比数量重要。100 张涵盖不同场景的图片比 1000 张同一场景的图片效果好得多。另外校准图片最好经过和推理时相同的预处理resize、归一化这样统计出来的激活值分布才准确。6.3 后处理对不上anchor顺序的陷阱ONNX 输出的 loc 和 conf 的排列顺序和 anchor 的生成顺序必须严格对应。我有一次自己写 anchor 生成代码时把特征层的顺序搞反了先生成 20×20 的再生成 80×80 的结果解码出来的框全部错位。调试了半天才发现是顺序问题。正确的顺序是按照 FPN 输出的顺序从大尺度80×80到小尺度20×20。每个尺度内先遍历 y 再遍历 x每个位置生成 2 个 anchor。这个顺序要和训练时保持一致否则解码结果完全不对。验证方法很简单用一张已知人脸位置的图片跑一遍推理看看解码后的框是不是落在人脸附近。如果框全部偏移或者尺寸不对大概率是 anchor 顺序问题。6.4 板端内存不足模型和图片的显存管理RK3588 的 NPU 有独立的内存区域但大小有限。如果模型太大或者同时加载多个模型可能会报out of memory。RetinaFace ResNet50 的 RKNN 模型大概 30-40 MB加上输入输出 tensor总共占 100 MB 左右一般没问题。但如果同时跑人脸检测和人脸识别两个模型就要注意内存分配了。解决办法用rknn_lite.init_runtime(core_mask...)时指定不同的核心让模型分散到不同 NPU 核心上。推理完成后及时rknn_lite.release()释放资源。如果板子内存实在紧张可以考虑用 MobileNet0.25 的轻量版模型。6.5 实际部署中的帧率稳定性问题在实验室跑单张图片时推理时间很稳定。但实际部署到视频流上帧率会波动。原因通常是图片解码耗时从摄像头读到的 MJPEG 或 H.264 流需要解码这部分在 CPU 上跑可能成为瓶颈。可以用 RK3588 的 VPU 硬件解码MPP来加速。内存拷贝OpenCV 的 Mat 和 NPU 的输入 tensor 之间需要拷贝如果图片尺寸大拷贝耗时不可忽略。可以用 RGA 做零拷贝的 resize 和格式转换。CPU 调度如果板子上同时跑其他任务CPU 抢占会导致后处理变慢。可以用taskset把推理进程绑定到特定核心。我的做法是把预处理和后处理都放到独立线程里和 NPU 推理流水线化。这样 NPU 在跑当前帧的时候CPU 已经在准备下一帧的数据了。实测可以把整体帧率提升 20%-30%。7. 一些让部署更顺手的经验RetinaFace 在 RK3588 上的部署核心难点其实不在 NPU 推理本身而在模型转换和后处理的细节。NPU 推理有官方工具链兜底只要版本匹配、算子支持基本不会出大问题。真正耗时间的是 ONNX 导出的算子兼容性、量化校准集的选择、anchor 顺序的对齐、NMS 的性能优化这些环节。我的建议是先在 PC 端用 onnxruntime 把整个推理流程跑通包括预处理、模型推理、后处理、NMS确保输出结果和 PyTorch 原版一致。然后再转 RKNN转完之后在 PC 端用 RKNN-Toolkit2 的仿真功能验证一遍。最后再上板子。这样每一步都有对照出了问题容易定位。另外RKNN-Toolkit2 的accuracy_analysis功能一定要用。它生成的报告能告诉你每一层的量化误差比盲目调参高效得多。如果某些层误差大优先考虑增加校准图片而不是急着上混合量化。板端部署时记得用verboseTrue打开 RKNNLite 的日志能看到每一层在哪个核心上跑、耗时多少。这些信息对性能调优很有帮助。如果发现某个层特别慢可能是被 fallback 到 CPU 了需要检查该层的算子是否被 NPU 支持。最后分享一个小技巧RK3588 的 NPU 支持多模型并行如果你的人脸检测和人脸识别是两个独立模型可以分别 init 到不同的 NPU 核心上用多线程同时推理。这样整体吞吐量能提升接近一倍。不过要注意两个模型之间的数据传递检测到的人脸框要裁剪出来送给识别模型这个裁剪和 resize 的操作最好用 RGA 做比 OpenCV 快很多。
返回列表