RK3588部署PyTorch模型：ONNX转RKNN全流程与避坑指南

1. 为什么要在RK3588上折腾PyTorch转RKNN这件事

手里有一块RK3588的板子，跑通了Ubuntu系统，连上了摄像头，然后想把训练好的PyTorch模型塞进去跑推理——这个流程听起来顺理成章，但真正动手的人都知道，从.pt文件到板子上流畅出结果，中间隔着的不是一条直线，而是一片雷区。我自己前前后后在这上面耗了将近三周，踩过的坑从环境版本冲突到算子不支持，从量化精度崩盘到内存对齐报错，几乎把能遇到的错误都遇了一遍。

这篇内容就是把我这段时间的实战经验完整梳理出来。核心关键词就几个：RK3588、PyTorch、RKNN、模型转换、部署。我会从模型训练完之后的导出环节讲起，一路说到板子上的推理验证，把每一步的操作意图、参数选择依据、以及最容易翻车的地方都摊开来讲。适合谁看？如果你手头有RK3588开发板，或者正在评估用RK3588做边缘推理方案，又或者你只是好奇PyTorch模型怎么变成NPU能吃的格式，这篇都能给你省下大量试错时间。

先说一个基本认知：RK3588的NPU不是万能加速器，它有自己的算子支持列表、有自己的量化偏好、有自己的内存管理逻辑。PyTorch模型直接扔进去是跑不起来的，必须经过RKNN-Toolkit2转换成.rknn格式。这个转换过程本质上是一次“翻译+压缩”，翻译是把PyTorch的算子映射到RKNN的算子，压缩是通过量化把FP32的权重压成INT8。翻译可能遇到“这个词字典里没有”的情况，压缩可能遇到“压完之后意思变了”的情况。我后面会逐一拆解怎么应对。

另外需要提前说明的是，RKNN-Toolkit2的版本和RK3588的NPU驱动版本之间存在严格的对应关系。我一开始用最新版的Toolkit去连老固件的板子，直接连不上，报错信息还很模糊。后来查了文档才知道，Toolkit2的1.5.x和1.6.x对应的NPU驱动版本不同，板子上的librknnrt.so版本必须和Toolkit匹配。这个点我会在环境搭建部分详细展开。

2. 环境搭建：PC端转换环境和板端运行环境要分开搞

2.1 PC端RKNN-Toolkit2的安装与版本选择

PC端的任务是把PyTorch模型转成RKNN格式，所以需要装RKNN-Toolkit2。这里第一个坑就是Python版本。RKNN-Toolkit2目前对Python 3.6、3.7、3.8、3.9、3.10都有对应的whl包，但并不是所有版本都同时支持所有功能。我实测下来，Python 3.8是最稳的，3.10在某些算子转换时会报一些奇怪的错误。如果你用conda管理环境，建议直接创建一个干净的3.8环境：

conda create -n rknn python=3.8 conda activate rknn

然后去官方仓库下载对应版本的whl包。注意，RKNN-Toolkit2的whl包分两种：一种是带cp38字样的，对应Python 3.8；另一种是cp310，对应3.10。下载错了装不上。安装命令就是普通的pip install，但要注意依赖的版本。我遇到过numpy版本过高导致Toolkit导入失败的情况，后来锁定在numpy==1.21.6就正常了。

pip install rknn_toolkit2-1.6.0-cp38-cp38-linux_x86_64.whl pip install numpy==1.21.6

装完之后用python -c "from rknn.api import RKNN"验证一下，不报错就说明PC端环境OK了。这里有个细节：如果你之前装过onnx或者torch，可能会和Toolkit的依赖产生冲突。我的建议是单独建一个conda环境专门做转换，不要和训练环境混在一起。训练环境里torch版本可能很新，但Toolkit依赖的onnx版本可能偏旧，混在一起迟早出问题。

2.2 板端NPU驱动与运行时库的匹配

板子这边需要确保NPU驱动和运行时库librknnrt.so的版本与PC端Toolkit匹配。查看板端版本的方法：

cat /sys/kernel/debug/rknpu/version

这个命令会输出NPU驱动的版本号。然后librknnrt.so通常在/usr/lib/目录下，可以用strings命令查看它的版本信息。关键是：Toolkit2 1.6.0要求板端librknnrt.so版本不低于1.6.0。如果板子固件比较老，需要单独更新这个so文件。更新方法很简单，把新版的librknnrt.so拷贝到板子的/usr/lib/目录，替换原来的文件，然后重启。

但这里有个隐藏坑：有些板子厂商的固件里，librknnrt.so是放在/usr/lib/aarch64-linux-gnu/下面的，而且可能有多个副本。你需要用find / -name "librknnrt.so"把所有副本都找出来，全部替换掉，否则运行时可能加载到旧版本。我就因为漏了一个副本，调试了半天才发现版本没生效。

另外，板端还需要安装rknn_server，这是PC端通过adb连接板子进行推理验证时用的服务。在板子上运行：

rknn_server &

如果没有这个命令，说明固件里没带，需要从SDK里编译或者找厂商要预编译版本。这个服务的作用是让PC端的Toolkit可以通过网络把模型推到板子上跑，方便调试。正式部署时不需要它，直接写C++或Python程序调用librknnrt.so就行。

2.3 PyTorch训练环境的注意事项

虽然转换是在PC端单独环境做的，但训练环境也有几个点会影响后续转换。首先是模型导出格式。PyTorch模型转RKNN有两条路：一条是torch.onnx.export导出ONNX，然后ONNX转RKNN；另一条是直接用Toolkit的load_pytorch接口加载.pt文件。我强烈建议走ONNX这条路，因为ONNX的算子表达更规范，Toolkit对ONNX的支持也更成熟。直接加载.pt文件经常遇到算子不支持或者图结构解析错误的问题。

导出ONNX时，opset_version建议设为11或12。设太高（比如15、16）Toolkit可能不认，设太低（比如9）某些算子表达不了。我实测opset_version=11兼容性最好。另外，导出时一定要指定input_names和output_names，并且确保输入尺寸是固定的。动态尺寸在RKNN上支持有限，后面会讲到。

还有一个容易忽略的点：PyTorch模型里的某些操作在导出ONNX时会被拆成多个算子，比如nn.Upsample可能被拆成Resize加Cast，这些拆出来的算子如果RKNN不支持，就会导致转换失败。所以训练时尽量用简单的算子组合，避免花哨的自定义层。如果必须用，后面我会讲怎么在RKNN里做自定义算子映射。

3. 模型转换全流程：从ONNX到RKNN的每一步操作

3.1 ONNX导出时的参数陷阱与验证方法

导出ONNX的代码看起来很简单，但参数没设对后面全是坑。先看一个标准的导出模板：

import torch import torch.onnx model = YourModel() model.load_state_dict(torch.load("model.pth")) model.eval() dummy_input = torch.randn(1, 3, 640, 640) torch.onnx.export( model, dummy_input, "model.onnx", export_params=True, opset_version=11, do_constant_folding=True, input_names=["input"], output_names=["output"], dynamic_axes=None )

这里有几个关键点。do_constant_folding=True会把一些常量计算提前算好，减小模型体积，但有时候会引入一些奇怪的算子，如果转换时报错，可以试着关掉它。dynamic_axes=None表示输入尺寸固定，这是RKNN最友好的情况。如果你确实需要动态batch，可以设dynamic_axes={"input": {0: "batch"}}，但RKNN对动态batch的支持需要额外配置，后面会讲。

导出之后必须验证ONNX模型是否正确。用onnxruntime跑一遍，和PyTorch的输出对比：

import onnxruntime as ort import numpy as np sess = ort.InferenceSession("model.onnx") onnx_out = sess.run(None, {"input": dummy_input.numpy()}) torch_out = model(dummy_input).detach().numpy() print(np.max(np.abs(onnx_out[0] - torch_out)))

如果误差在1e-5以内，说明导出没问题。如果误差很大，说明某些算子导出时行为不一致，需要排查。我遇到过一次nn.Hardswish在opset 11下导出后计算结果偏差很大，后来换成nn.ReLU就正常了。所以导出后的数值验证这一步绝对不能省。

3.2 RKNN转换配置：mean_values和std_values的正确设置

ONNX转RKNN的核心是RKNN.config()这个接口。参数很多，但最关键的几个是mean_values、std_values、target_platform和quantized_dtype。

from rknn.api import RKNN rknn = RKNN(verbose=True) rknn.config( mean_values=[[0, 0, 0]], std_values=[[255, 255, 255]], target_platform="rk3588", quantized_dtype="asymmetric_quantized-8", optimization_level=3 )

mean_values和std_values是归一化参数。这里有个大坑：RKNN的归一化是在NPU内部做的，所以你的模型在训练时如果已经做了归一化，这里就要设成0和1，否则会重复归一化。我一开始没注意，训练时用了transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])，转换时又设了mean和std，结果推理输出完全不对。后来把mean设成0、std设成1，让NPU不做额外处理，问题就解决了。

target_platform必须设为rk3588，设成rk3568或其他的会导致生成的模型在RK3588上跑不了。quantized_dtype选asymmetric_quantized-8是INT8量化，精度和速度的平衡最好。如果精度要求极高，可以用float16，但速度会慢一些，而且RK3588的NPU对FP16的支持不如INT8充分。

optimization_level设3会做更多的图优化，比如算子融合、常量折叠等，通常能提升推理速度。但有时候优化过头会导致精度下降，如果发现转换后精度不对，可以降到2或1试试。

3.3 量化数据集准备：为什么不能随便找几张图糊弄

INT8量化需要一个校准数据集，用来统计激活值的分布，确定量化参数。这个数据集的质量直接决定量化后的精度。我见过有人随便拿几张纯色图做校准，结果量化后模型输出全是乱的。

校准数据集的要求：至少100张图，最好300-500张，要覆盖实际推理时可能遇到的各种场景。比如你做目标检测，校准集里就要有白天、夜晚、室内、室外、不同光照条件的图。图片尺寸要和模型输入尺寸一致，或者至少保持相同的宽高比。

rknn.load_onnx(model="model.onnx") rknn.build(do_quantization=True, dataset="calibration_dataset.txt")

calibration_dataset.txt是一个文本文件，每行是一张图片的路径。注意路径要写绝对路径，相对路径有时候会找不到。另外，图片格式建议用JPG或PNG，不要用BMP，Toolkit对BMP的支持不太稳定。

如果量化后精度下降太多，可以尝试混合量化：把某些对精度敏感的层保持FP16，其他层用INT8。Toolkit支持通过hybrid_quantization配置来实现，但操作比较复杂，需要逐层分析。我的经验是，大部分常见模型（ResNet、YOLO、MobileNet）纯INT8量化后精度损失在1%以内，不需要混合量化。只有一些特殊结构（比如Transformer里的Attention）可能需要特殊处理。

3.4 转换结果验证：PC端模拟推理与精度对比

转换完成后，不要急着推到板子上，先在PC端用Toolkit的模拟推理功能验证一下：

rknn.init_runtime() outputs = rknn.inference(inputs=[dummy_input.numpy()])

把outputs和ONNX的输出对比，看误差有多大。如果误差在可接受范围内（比如分类任务Top-1一致，检测任务mAP下降不超过2%），就可以进行下一步。如果误差很大，说明量化有问题，需要调整校准集或者换量化方式。

这里有个细节：rknn.inference的输入格式是NHWC还是NCHW？RKNN默认是NHWC，但PyTorch是NCHW。所以如果你直接传NCHW的numpy数组进去，结果会不对。需要在config里设置data_format="nchw"，或者在传数据前手动transpose。我建议在config里设好，省得后面搞混。

4. 板端部署实战：从模型推送到推理验证

4.1 adb连接与模型推送的完整操作

PC端验证通过后，把.rknn文件推到板子上。用adb连接：

adb connect 192.168.1.100:5555 adb push model.rknn /data/

IP地址换成你板子的实际地址。如果adb连不上，先检查板子的adb服务有没有开，有些固件默认不开adb。可以在板子上运行adbd &手动启动。另外，USB连接和网络连接都可以，USB更稳定但需要驱动，网络更方便但受网速影响。

推完模型后，在板子上写一个简单的Python推理脚本：

from rknnlite.api import RKNNLite rknn = RKNNLite() ret = rknn.load_rknn("model.rknn") ret = rknn.init_runtime() import cv2 import numpy as np img = cv2.imread("test.jpg") img = cv2.resize(img, (640, 640)) img = img[:, :, ::-1] # BGR to RGB img = np.expand_dims(img, axis=0) outputs = rknn.inference(inputs=[img]) print(outputs[0].shape)

注意板端用的是RKNNLite，不是RKNN。RKNNLite是轻量级运行时，只支持推理，不支持转换。init_runtime()的时候可以指定core_mask来绑定NPU核心，RK3588有三个NPU核心，可以并行跑多个模型。

4.2 性能调优：NPU核心绑定与内存优化

RK3588的NPU有三个核心，默认是自动调度。如果你的应用需要同时跑多个模型，可以手动绑定核心：

rknn.init_runtime(core_mask=RKNNLite.NPU_CORE_0)

NPU_CORE_0、NPU_CORE_1、NPU_CORE_2分别对应三个核心，NPU_CORE_0_1_2是三个一起用。对于单个模型，用三个核心一起跑通常最快，但功耗也最高。如果是多模型并行，每个模型绑一个核心，整体吞吐量更大。

内存方面，RKNN模型加载后会占用NPU的专用内存，不占CPU内存。但如果输入图片很大，预处理阶段会占CPU内存。建议在板子上用cv2.resize时指定interpolation=cv2.INTER_LINEAR，比默认的INTER_NEAREST快很多。另外，如果板子内存紧张，可以在init_runtime时设置rknn.init_runtime(perf_mode=True)来启用性能模式，会优化内存分配。

4.3 精度对齐：板端输出与PC端不一致的排查

有时候PC端模拟推理结果正常，但板端跑出来不一样。常见原因有三个：一是板端librknnrt.so版本和PC端Toolkit不匹配，这个前面说过；二是输入数据的预处理不一致，比如PC端用了归一化但板端没做；三是NPU的量化计算和PC端模拟的量化计算有细微差异，这个通常误差很小，不影响结果。

排查方法：把同一张图分别用PC端和板端跑，输出保存下来对比。如果误差在1e-3以内，基本正常。如果误差很大，先检查预处理代码是否完全一致，再检查版本匹配。我遇到过一次板端输出全零的情况，后来发现是输入图片的通道顺序搞反了，PC端是RGB，板端传了BGR，导致NPU内部计算出错。

5. 常见问题速查与避坑经验汇总

5.1 转换阶段报错排查表

5.2 推理阶段常见异常与处理

板端推理时最常见的异常是Segmentation fault，通常是因为librknnrt.so版本不对或者模型文件损坏。可以先重新推送模型文件，如果还不行就检查so版本。另一个常见问题是推理速度远低于预期，这时候要检查core_mask是否设置正确，以及输入图片的预处理是否成了瓶颈。我实测下来，YOLOv8s在RK3588上跑640x640的输入，INT8量化后单核推理大约30ms，三核并行大约12ms。如果你的速度差很多，大概率是预处理拖了后腿。

还有一个隐蔽的坑：RKNN模型在多次推理后可能会出现内存泄漏。这是因为每次inference都会分配新的输出buffer，如果不手动释放，内存会持续增长。解决方法是在循环推理时复用输入输出buffer，或者定期重启推理进程。Toolkit的文档里没怎么提这一点，但我在长时间跑视频流时遇到了，后来改成每处理1000帧重启一次推理服务，问题就消失了。

5.3 我的个人经验总结

最后分享几个我觉得最有用的经验。第一，永远先在PC端把ONNX和RKNN的精度验证做完再上板子，板子调试比PC端麻烦十倍。第二，校准集一定要用真实场景的图，不要用训练集的子集，因为训练集可能已经经过增强，分布和实际推理时不一样。第三，版本管理要严格，Toolkit、驱动、so文件三者版本必须一致，建议在项目开始时就锁定版本，不要中途升级。第四，输入尺寸尽量用64的倍数，比如640、512、384，这样NPU的内存对齐效率最高，速度会快一些。

这个流程后续还可以扩展的方向包括：多模型并行推理的调度策略、视频流场景下的零拷贝传输、以及自定义算子的注册与编译。如果后面有机会，我再把这些内容整理出来。