从PyTorch到RK3588：高效模型转换与NPU推理实战指南

关注

VIP免费原创最新推荐文章于 2026-10-15 10:16:22 发布 274 阅读 AI 写同款· GEO 优化 ›

如何在RK3588开发板上用rknntoolkit2快速部署PyTorch模型（附完整代码）本文详细介绍了在RK3588开发板上使用rknntoolkit2快速部署PyTorch模型的全流程，包括环境配置、模型转换、量化优化和高效推理实现。通过完整代码示例，帮助开发者掌握从pt模型到RKNN格式的转换技巧，并充分利用RK3588 NPU的6TOPS算力进行边缘计算部署。阅读详情

1. 为什么要把PyTorch模型部署到RK3588的NPU上？

如果你正在做嵌入式AI项目，比如智能摄像头、机器人或者边缘计算盒子，肯定遇到过这样的问题：在PC上训练好的PyTorch模型跑得飞快，但一到嵌入式设备上就卡成PPT。我之前在一个安防项目里就踩过这个坑，用ResNet50做实时视频分析，在PC上能跑到100多帧，但放到普通的ARM开发板上，一帧就要处理几百毫秒，完全没法用。

这时候RK3588的NPU（神经网络处理单元）就是你的救星。这颗芯片内置的NPU算力最高能达到6TOPS，什么概念呢？就是每秒能进行6万亿次运算。对比一下，用CPU跑同样的模型可能只有0.1-0.5TOPS，性能直接提升几十倍。而且NPU是专门为神经网络计算设计的，功耗还特别低，非常适合需要长时间运行的边缘设备。

但问题来了，PyTorch模型不能直接在NPU上跑，需要经过转换。这个转换过程就像把Word文档转成PDF——内容不变，但格式变了，适配不同的“阅读器”。整个过程分为三步：PyTorch转ONNX、ONNX转RKNN、最后在板子上用NPU推理。听起来有点复杂？别担心，我一步步带你走通，把我踩过的坑都告诉你。

2. 环境搭建：别在第一步就翻车

2.1 PC端环境准备

我强烈建议用Ubuntu 20.04或22.04，这是官方测试最多的系统版本。Windows和MacOS虽然理论上也能用，但兼容性问题会让你头疼死。我试过在Windows的WSL2里搞，各种库版本冲突，最后老老实实换回了Ubuntu。

首先安装必要的系统依赖：

sudo apt-get update
sudo apt-get install python3.8 python3.8-dev python3-pip
sudo apt-get install libxslt1-dev zlib1g zlib1g-dev libglib2.0-0 libsm6 libgl1-mesa-glx libprotobuf-dev gcc

重要提醒：zlib1g中间是数字1，不是字母l，我当初就因为这个拼写错误折腾了半天。

2.2 创建Python虚拟环境

直接用系统Python很容易把环境搞乱，用conda或者virtualenv隔离一下最稳妥。我习惯用conda，因为管理不同Python版本更方便：

# 创建名为rknn的虚拟环境，指定Python 3.8
conda create -n rknn python=3.8
conda activate rknn

为什么选Python 3.8？因为RKNN Toolkit2对3.6和3.8支持最好，3.9、3.10可能会有奇怪的兼容性问题。我试过3.10，有些依赖包装不上，最后还是退回3.8。

2.3 安装RKNN Toolkit2

先去GitHub下载最新版的RKNN Toolkit2：

git clone https://github.com/rockchip-linux/rknn-toolkit2.git
cd rknn-toolkit2

安装Python依赖时有个大坑：官方给的requirements文件里有些版本已经过时了。我整理了一个实测可用的版本列表：

pip install numpy==1.19.5
pip install onnx==1.12.0
pip install onnxruntime==1.12.1
pip install torch==1.10.0
pip install torchvision==0.11.1
pip install opencv-python==4.5.5.64
pip install protobuf==3.20.3

注意protobuf的版本，3.20.0以上会有问题，我用的3.20.3是稳定的。如果安装时提示某个包版本冲突，先卸载冲突的包再重装。

然后安装RKNN Toolkit2本体：

cd packages
# 根据你的Python版本和系统架构选择对应的whl文件
pip install rknn_toolkit2-1.5.2+b642f30c-cp38-cp38m-linux_x86_64.whl

验证安装是否成功：

python -c "from rknn.api import RKNN; print('导入成功')"

如果没报错，PC端环境就准备好了。

3. 模型转换：从PyTorch到RKNN的完整流程

3.1 第一步：PyTorch转ONNX

ONNX（Open Neural Network Exchange）是个中间格式，相当于AI模型的“普通话”，各种框架都能听懂。转换时要注意几个关键点：

import torch
import torch.onnx

# 加载你的PyTorch模型
model = YourModel()
model.load_state_dict(torch.load('your_model.pth'))
model.eval()  # 一定要切换到eval模式！

# 准备一个示例输入
dummy_input = torch.randn(1, 3, 224, 224)  # [batch, channels, height, width]

# 设置输入输出名称
input_names = ['input']
output_names = ['output']

# 导出ONNX
torch.onnx.export(
    model,
    dummy_input,
    'model.onnx',
    export_params=True,
    opset_version=11,  # 重要：opset版本不能太高
    do_constant_folding=True,
    input_names=input_names,
    output_names=output_names,
    dynamic_axes={'input': {0: 'batch_size'}} if dynamic else None
)

我踩过的坑：

opset版本：RKNN Toolkit2对opset 11支持最好，12以上可能会有算子不支持。如果你导出时用了高版本，转换RKNN时会报错。
动态维度：如果想支持可变batch size或可变尺寸，要设置dynamic_axes。但注意，动态维度会增加转换复杂度，初次尝试建议先用固定尺寸。
自定义算子：如果你的模型用了非标准算子，比如某些特殊的激活函数，需要先实现对应的ONNX算子或者用标准算子替换。

转换完成后一定要验证ONNX模型是否正确：

import onnxruntime as ort
import numpy as np

# 加载ONNX模型
ort_session = ort.InferenceSession('model.onnx')

# 准备相同输入
input_data = np.random.randn(1, 3, 224, 224).astype(np.float32)

# ONNX推理
onnx_outputs = ort_session.run(None, {'input': input_data})

# PyTorch推理
with torch.no_grad():
    torch_outputs = model(torch.from_numpy(input_data))

# 对比结果
print(f"最大误差: {np.max(np.abs(onnx_outputs[0] - torch_outputs.numpy()))}")

误差应该在1e-5以内，如果差太多，说明转换有问题。

3.2 第二步：ONNX转RKNN

这是最关键的一步，直接决定了模型在NPU上的性能。先看基础转换代码：

from rknn.api import RKNN

def onnx_to_rknn(onnx_path, rknn_path, target_platform='rk3588'):
    # 创建RKNN对象
    rknn = RKNN(verbose=True)  # verbose=True可以看到详细日志
    
    # 配置模型参数
    print('--> Config model')
    rknn.config(
        mean_values=[[127.5, 127.5, 127.5]],  # 归一化参数
        std_values=[[127.5, 127.5, 127.5]],
        quant_img_RGB2BGR=False,  # 输入图像是RGB还是BGR
        target_platform=target_platform,
        optimization_level=3,  # 优化等级，3是最高
        quantized_algorithm='normal',
        quantized_method='channel'
    )
    
    # 加载ONNX模型
    print('--> Loading model')
    ret = rknn.load_on