多彩编程 多彩编程MZPH · CODE BLOG
ARTICLE DETAIL

文章详情

深耕前端与后端开发技术的一线实战笔记与踩坑复盘。

RK3588 NPU部署RetinaFace:PyTorch转RKNN全流程与性能调优

RK3588 NPU部署RetinaFace:PyTorch转RKNN全流程与性能调优 1. 为什么要在RK3588上跑RetinaFaceRK3588这颗芯片做边缘AI的人应该都不陌生。8核CPU4×A764×A55、Mali-G610 GPU、6TOPS算力的NPU支持INT4/INT8/INT16混合量化视频编解码能力也够强基本上是目前国产边缘计算平台里性价比最能打的一档。我手头这块板子跑了大半年的各种模型从YOLO系列到各种Backbone踩过的坑攒了一箩筐。RetinaFace这个模型做人脸检测的朋友肯定都听过。它在WiderFace榜单上表现一直很稳单阶段检测能同时输出人脸框和5个关键点双眼、鼻尖、左右嘴角对于人脸对齐、人脸属性分析这类下游任务来说关键点信息非常实用。相比MTCNN那种三级级联的结构RetinaFace一步到位速度更快精度也不差。但问题来了——PyTorch训练出来的模型直接往RK3588上扔是跑不起来的。RK3588的NPU只认RKNN格式你需要把PyTorch的.pth或者.pt文件先转成ONNX再通过RKNN-Toolkit2转成.rknn最后在板子上用RKNPU2的运行时库加载推理。这条链路听起来简单实际上每一步都有坑ONNX导出时的算子兼容性、RKNN转换时的量化校准、板端推理时的内存对齐、输入输出的layout差异……任何一个环节出问题你拿到的就是一堆乱码或者直接报错。这篇内容就是把我从PyTorch到RKNN的完整流程梳理一遍包括环境搭建、模型转换、量化校准、板端部署、性能调优以及那些文档里不会写的坑。适合手里有RK3588板子、想在NPU上跑人脸检测的开发者也适合正在做模型端侧部署、对RKNN工具链还不熟的朋友。整个流程我会尽量给出可直接复现的命令和代码你照着走一遍基本能跑通。2. 环境搭建与工具链选型2.1 PC端环境Ubuntu Conda RKNN-Toolkit2RKNN-Toolkit2是瑞芯微官方提供的模型转换工具只能在Linux x86_64上跑Windows和Mac都不行。我试过在WSL2里装能用但USB设备直通有时候会抽风建议还是用原生Ubuntu 20.04或22.04。我目前用的是Ubuntu 22.04Python 3.8的Conda环境这个组合最稳。先建环境conda create -n rknn python3.8 conda activate rknn然后装PyTorch。注意RKNN-Toolkit2对PyTorch版本有要求太新的版本可能不兼容。我实测下来PyTorch 1.13.1 torchvision 0.14.1这个组合最稳CUDA版本用11.7就行pip install torch1.13.1cu117 torchvision0.14.1cu117 -f https://download.pytorch.org/whl/torch_stable.html接下来装RKNN-Toolkit2。从瑞芯微的GitHub仓库或者官方开发者网站下载whl包注意版本要和板子上的RKNPU2运行时版本匹配。我用的rknn_toolkit2-1.6.0对应的RKNPU2是2.0.0pip install rknn_toolkit2-1.6.0-cp38-cp38-linux_x86_64.whl装完之后验证一下from rknn.api import RKNN print(RKNN().version)能打印出版本号就说明装好了。如果报错说找不到librknnrt.so那是动态库路径的问题把RKNN-Toolkit2的lib目录加到LD_LIBRARY_PATH里就行。注意RKNN-Toolkit2和RKNN-Toolkit1是两套东西RK3588只能用Toolkit2别装错了。另外Toolkit2的版本和板端RKNPU2的版本必须匹配否则转换出来的模型加载会失败。2.2 板端环境RKNPU2 OpenCV板子这边我用的固件是官方Ubuntu 20.04的镜像内核版本5.10。RKNPU2的运行时库一般固件里已经带了如果没有可以从官方仓库编译安装。检查一下ls /usr/lib/librknnrt.so有的话就不用管了。OpenCV板子上一般也预装了但版本可能比较老建议自己编译一个带Python绑定的版本或者直接用pip装opencv-python-headlesspip install opencv-python-headless numpy板端的Python版本也要注意RKNPU2的Python绑定目前支持3.6到3.9我用的是3.8没问题。2.3 模型准备RetinaFace的PyTorch实现RetinaFace有很多开源实现我选的是biubug6的PyTorch版本结构清晰导出ONNX比较方便。下载下来之后你需要一个训练好的权重文件或者自己训练一个。我这里用的是ResNet50 backbone的版本输入尺寸640×640精度和速度比较平衡。先加载模型确认能正常推理import torch from models.retinaface import RetinaFace net RetinaFace(cfgcfg, phasetest) net.load_state_dict(torch.load(RetinaFace_Resnet50.pth, map_locationcpu)) net.eval()跑一张测试图确认输出正常。这一步很重要如果PyTorch这边就有问题后面转换肯定也会出问题。3. 从PyTorch到ONNX导出细节与算子兼容性3.1 ONNX导出步骤ONNX导出是整个链路的第一步也是最容易出问题的一步。RetinaFace的输出结构比较特殊它有多个输出分支三个尺度的分类输出、三个尺度的回归输出、三个尺度的关键点输出再加上一个额外的输出。导出的时候要把这些输出都保留下来。import torch.onnx dummy_input torch.randn(1, 3, 640, 640) torch.onnx.export( net, dummy_input, retinaface.onnx, opset_version11, input_names[input], output_names[cls_8, cls_16, cls_32, reg_8, reg_16, reg_32, kps_8, kps_16, kps_32, conf], dynamic_axes{input: {0: batch}} )opset_version选11这个版本RKNN-Toolkit2支持得最好。dynamic_axes把batch维度设成动态的方便后面做batch推理。但要注意RKNN对动态batch的支持有限实际部署时通常还是固定batch size。3.2 常见导出错误与解决导出过程中最常见的报错是算子不支持。RetinaFace里用到了F.interpolate做上采样如果opset版本太低会报错说不支持resize。解决办法就是升opset到11以上。另外如果模型里有自定义算子需要自己写ONNX的符号函数。还有一个坑是输出顺序。PyTorch的forward返回的是一个列表导出到ONNX之后输出的顺序可能和你想象的不一样。建议导出之后用Netron打开看一下确认每个输出的名字和形状对不对。pip install netron netron retinaface.onnxNetron里能看到每个节点的输入输出形状检查一下有没有异常。如果发现某个输出形状不对大概率是forward里的reshape或者permute操作导致的需要在导出前调整。3.3 ONNX模型验证导出之后用onnxruntime跑一遍和PyTorch的输出对比import onnxruntime as ort import numpy as np sess ort.InferenceSession(retinaface.onnx) input_name sess.get_inputs()[0].name outputs sess.run(None, {input_name: dummy_input.numpy()}) # 和PyTorch输出对比 with torch.no_grad(): torch_outputs net(dummy_input) for i, (onnx_out, torch_out) in enumerate(zip(outputs, torch_outputs)): diff np.abs(onnx_out - torch_out.numpy()).max() print(fOutput {i}: max diff {diff})如果max diff在1e-5以内说明导出没问题。如果差很多检查一下是不是某个算子被替换了或者输入预处理不一致。实操心得ONNX导出时建议把模型里的后处理比如decode、NMS剥离出来只导出纯网络部分。后处理放在板端用C或Python实现这样转换更稳定也方便调试。4. RKNN模型转换与量化校准4.1 RKNN转换脚本编写ONNX有了接下来用RKNN-Toolkit2转成.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, 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)mean_values和std_values要和训练时的预处理一致。RetinaFace官方实现里用的是mean[104, 117, 123]std[1, 1, 1]这里保持一致。4.2 量化校准数据集准备do_quantizationTrue的时候需要提供一个校准数据集。dataset.txt里每行是一张图片的路径图片数量建议在100到500张之间太少量化误差大太多转换时间太长。图片要覆盖各种场景不同光照、不同角度、不同人脸大小。find ./calib_images -name *.jpg dataset.txt校准图片的预处理要和推理时一致包括resize到640×640、归一化等。RKNN-Toolkit2会自动读取图片并做预处理但前提是你的mean_values和std_values配置正确。4.3 量化精度调优量化之后精度通常会掉一点这是正常的。如果掉太多可以尝试以下方法增加校准图片数量特别是增加难样本小脸、遮挡、暗光调整optimization_level从3降到2或1转换会更保守精度可能更好对敏感层使用混合量化RKNN-Toolkit2支持通过hybrid_quantization_step1和step2做混合量化我实测下来ResNet50 backbone的RetinaFaceINT8量化之后WiderFace的AP大概掉1到2个百分点基本可以接受。如果对精度要求极高可以考虑用FP16但速度会慢一些。注意量化校准的时候如果某张图片读取失败RKNN-Toolkit2会直接报错退出。建议先用脚本检查一遍dataset.txt里的路径是否都存在图片是否能正常解码。5. 板端部署与推理实现5.1 RKNPU2 Python API推理板子上用RKNPU2的Python接口加载模型from rknnlite.api import RKNNLite import cv2 import numpy as np rknn RKNNLite() ret rknn.load_rknn(retinaface.rknn) ret rknn.init_runtime(core_maskRKNNLite.NPU_CORE_0_1_2) img cv2.imread(test.jpg) img cv2.resize(img, (640, 640)) img img.astype(np.float32) img img - np.array([104, 117, 123]) img img[np.newaxis, :, :, :] outputs rknn.inference(inputs[img])core_mask可以指定用哪些NPU核心RK3588有三个NPU核心可以单独用也可以组合用。多核并行能提升吞吐量但延迟不一定降低。5.2 后处理实现RKNN的输出是原始的特征图需要自己做decode和NMS。RetinaFace的decode逻辑不复杂主要是把anchor和预测的偏移量结合起来得到最终的框和关键点。def decode(loc, priors, variances): boxes np.concatenate(( priors[:, :2] loc[:, :2] * variances[0] * priors[:, 2:], priors[:, 2:] * np.exp(loc[:, 2:] * variances[1]) ), axis1) boxes[:, :2] - boxes[:, 2:] / 2 boxes[:, 2:] boxes[:, :2] return boxesNMS用OpenCV自带的就行indices cv2.dnn.NMSBoxes(boxes.tolist(), scores.tolist(), 0.5, 0.4)关键点的decode类似也是用偏移量乘以prior的宽高再加上prior的中心点。5.3 性能实测与调优我实测下来RK3588单核NPU跑RetinaFace 640×640单帧推理大概在25到30毫秒左右三核并行能降到15毫秒以内。加上前后处理整体单帧延迟在40毫秒左右差不多25FPS。这个性能做人脸检测足够了。如果速度不够可以尝试以下优化降低输入分辨率比如从640×640降到320×320速度能翻倍但小脸检测会变差减少backbone的层数用MobileNet替换ResNet50把后处理放到NPU上做但RKNN对NMS的支持有限需要自己实现实操心得板端推理时输入图片的预处理尽量用RGA硬件加速不要用CPU做resize和颜色空间转换。RGA是RK3588的2D加速器做resize和格式转换几乎不占CPU能省不少时间。6. 常见问题与排查技巧实录6.1 模型转换失败排查转换失败最常见的原因是算子不支持。RKNN-Toolkit2的日志里会打印哪个算子不支持你可以根据日志去查RKNN的算子支持列表。如果确实不支持有两个办法一是修改模型结构用支持的算子替换二是把不支持的算子放到CPU上跑RKNN支持自定义算子但性能会受影响。另一个常见问题是输入形状不匹配。ONNX的输入是动态batch但RKNN默认是固定batch。如果转换时报错说shape不匹配把dynamic_axes去掉固定batch size再试。6.2 板端推理报错排查板端加载模型时报错首先检查RKNPU2的版本和RKNN-Toolkit2的版本是否匹配。版本不匹配是最常见的原因。其次检查librknnrt.so的路径是否正确可以用ldd命令看依赖有没有缺失。推理时输出全零或者乱码大概率是输入预处理不对。检查mean和std是否和转换时一致输入图片的layout是NHWC还是NCHW。RKNN默认输入是NHWC如果你传的是NCHW需要先transpose。6.3 精度下降问题排查量化之后精度下降先确认校准数据集是否覆盖了实际场景。如果校准集和测试集分布差异大量化误差会很大。其次检查量化配置asymmetric_quantized-8通常比symmetric精度好但速度可能稍慢。如果精度还是不行可以尝试逐层量化分析找出误差最大的层对这一层使用FP16。RKNN-Toolkit2支持混合量化但操作比较复杂需要两步转换。问题现象可能原因解决方法转换时报算子不支持ONNX算子RKNN不支持替换算子或使用自定义算子板端加载模型失败RKNPU2版本不匹配升级或降级RKNPU2推理输出全零输入预处理错误检查mean/std和layout量化后精度掉太多校准集不具代表性增加难样本调整量化配置推理速度慢单核NPU预处理占CPU多核并行RGA加速预处理6.4 内存与性能问题RK3588板子内存一般4GB或8GB跑RetinaFace这种模型绰绰有余。但如果同时跑多个模型或者batch size开太大可能会OOM。建议推理时用固定batch size不要动态分配。性能方面NPU的频率是可以调的默认可能不是最高频。可以通过sysfs调整echo performance /sys/class/devfreq/fdab0000.npu/governor这样NPU会一直跑在最高频速度会快一些但功耗和发热会增加。注意调整NPU频率之前确认散热措施到位否则长时间高负载可能会触发降频反而影响稳定性。7. 一些个人体会这套流程我前前后后跑了不下十遍每次换模型或者换板子都要重新踩一遍坑。最大的感受是RKNN工具链的版本管理太重要了PC端的Toolkit2和板端的RKNPU2必须严格匹配否则各种莫名其妙的错误。建议在项目开始之前先把版本对齐记录下来后面换环境的时候直接照抄。另外量化校准这一步千万别偷懒。我见过太多人随便找几十张图就跑量化结果精度掉得没法看回头还要重新来。校准集的质量直接决定量化模型的质量花点时间准备几百张有代表性的图片绝对值得。最后板端推理的性能调优是个细活。RGA加速、NPU多核、频率调整这些手段叠加起来能把性能压榨到极限。但也要注意平衡不要为了追求极致速度牺牲稳定性。我在实际项目中通常会留20%的性能余量避免满负载运行导致的各种问题。
返回列表