GPU 训练 --> NPU 部署

• 模型训练：windows+pycharm+pytorch+gpu搭建训练的深度学习模型
• 部署：瑞芯微 RK3588 NPU Linux
• 路线：PyTorch(GPU)–>ONNX–>RKNN Toolkit–>model.rknn–>RK3588 NPU 推理

Step 0: 环境准备

1️⃣ Linux PC 端（ONNX转换为RKNN使用）
👉 RKNN Toolkit2 官方不支持 Windows 做模型转换
👉 一定在 x86 Linux PC 上做转换，不在板子上

• Ubuntu 18.04 / 20.04
• Python 3.8（官方最稳）
• RKNN Toolkit2

1. 虚拟机（Windows + WSL2）

# 以管理员身份打开命令提示符或PowerShell
wsl --install

参考链接：

https://blog.csdn.net/weixin_44818765/article/details/157736533
https://www.runoob.com/linux/windows-wsl-linux.html

2. Anaconda安装
   安装包可以从镜像源获取，打开网址https://mirrors.tuna.tsinghua.edu.cn/anaconda/archive/，选择对应的版本。 获取安装包后，在linux系统中使用默认shell执行下面命令安装：

./Anaconda3-2021.05-Linux-x86_64.sh

然后回车，一直往下滑动看完license，最后输入yes后，继续按下回车,然后进入安装，安装完成后输入yes初始化anaconda3。
重启系统或者使用命令source ~/.bashrc 后打开终端会进入默认anaconda环境里(命令行提示符前面有一个base)，就可以使用conda来创建虚拟环境和包管理：

# 重启或者使用命令source ~/.bashrc进入anaconda环境
source ~/.bashrc
# 下面列出一些常用conda命令，更多命令参考conda --help
# 创建一个环境，env_name是环境名称，后面可以指定python的版本等等
conda create -n env_name ....
# 列出创建的虚拟环境
conda info --envs
# 激活环境，不指定env_name时，默认进入base环境
conda activate env_name
# 移除环境，其中env_name是环境名称
conda remove -n env_name --all
# 退出当前conda环境
conda deactivate
# 每次开启终端都会默认进入anaconda环境，使用下面命令后重启就不会默认启动虚拟环境
conda config --set auto_activate_base false

默认Anaconda安装会带很多默认的库，要轻量版的可以使用 Miniconda ， 默认只包含了python和conda，通过pip和conda来安装所需要的包。可以点击 链接 去下载安装，安装和使用操作类似。
Anaconda配置使用国内镜像源（可选）：

# 查看默认的源
conda config --show channels

# 使用清华源，也可以自行添加其他第三方源等等
conda config --add channels https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/main
conda config --add channels https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/r

conda config --set custom_channels.auto https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud/

conda config --set show_channel_urls yes

3. rknn-toolkit2安装
   RKNN Toolkit2 开发套件运行在PC (x86_64/arm64)平台上，提供了模型转换、 量化功能、模型推理、性能和内存评估、量化精度分析、模型加密等功能。
   从GitHub 下载 wheel：

rknn-toolkit2/packages/x86_64/rknn_toolkit2-2.3.2-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl

pip install rknn_toolkit2-2.3.2-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl

测试，没报错即可：

python
>>> from rknn.api import RKNN
>>> rknn = RKNN()

4. Pycharm安装（可选）：
   下载安装包，linux的安装包是以.tar.gz格式的
   pycharm官方网站

tar zxf pycharm-2025.3.2.1.tar.gz 
cd pycharm-2025.3.2.1/bin
./pycharm.sh     ##运行pycharm安装脚本

Step 1: 导出 ONNX（Pytorch(GPU)–>ONNX）

👉在windows下即可进行转换

1. 生成 ONNX

import torch
import get_model

# 1️⃣ 创建模型
model = get_model()

# 2️⃣ 加载权重
state_dict = torch.load("./weights/mode1.pth", map_location="cpu")

model.load_state_dict(state_dict, strict=False)

# 3️⃣ 切换推理模式（非常重要）
model.eval()

dummy_input = torch.randn(1, 1, 8, 750)

# 导出ONNX
torch.onnx.export(
    model,                      # 模型
    dummy_input,                # 输入样例
    "model1.onnx",           # 输出文件
    export_params=True,         # ⭐ 把权重一起写入
    opset_version=11,           # ⭐ RK3588 推荐
    do_constant_folding=True,   # 常量折叠优化
    input_names=["input"],
    output_names=["output"],
    dynamic_axes=None           # ⭐ 禁止动态shape
)

2. 验证与优化ONNX
   安装：

pip install onnx onnxsim -i https://pypi.tuna.tsinghua.edu.cn/simple

检查：

import onnx

model = onnx.load("eeg_model.onnx")
onnx.checker.check_model(model)
print("ONNX OK")

优化：

python -m onnxsim model.onnx model_sim.onnx

3. 模型转ONNX注意事项

• 输入 shape 必须固定
• 不要用动态 batch
  可能的报错：

UnsupportedOperatorError:
Exporting operator ::_convolution_mode to ONNX opset 11 not supported

解决方法：

# 将模型中的
padding="same"
# 修改为确定的padding数，或下面的通用公式
padding = kernel_size // 2

小工具：使用Netron可查看.onnx模型结构

Step 2: 导出 RKNN（ONNX–>RKNN）

👉进行该步骤前，确保Step 0已经完成
👉Linux x86环境下进行
转换脚本：

⭐ convert.py

from rknn.api import RKNN

ONNX_MODEL = 'model.onnx'
RKNN_MODEL = 'model.rknn'

rknn = RKNN(verbose=True)

# ⭐⭐⭐ 稳定配置
rknn.config(
    target_platform='rk3588',
    optimization_level=3,
)

print('--> Loading ONNX')
ret = rknn.load_onnx(
    model=ONNX_MODEL,
    inputs=['input'],
    input_size_list=[[1,1,8,750]]
)
if ret != 0:
    exit(ret)

print('--> Building')
ret = rknn.build(do_quantization=False)
if ret != 0:
    exit(ret)

print('--> Export RKNN')
ret = rknn.export_rknn(RKNN_MODEL)
if ret != 0:
    exit(ret)

print('DONE')

运行：

python convert.py

注意事项：若报错，需追溯到模型搭建环节上。

Step 3: 部署到RK3588（在 RK3588 上运行 .rknn）

① 把模型.rknn拷贝到板子上
② 参考Step 0安装相关python环境
③ 板端安装 runtime
RK3588 上需要：

rknn-toolkit-lite2
从GitHub 下载 wheel：rknn-toolkit-lite2/packages/rknn_toolkit_lite2-2.3.2-cp38-cp38-manylinux_2_17_aarch64.manylinux2014_aarch64.whl

pip install rknn_toolkit_lite2-2.3.2-cp38-cp38-manylinux_2_17_aarch64.manylinux2014_aarch64.whl

④ 推理脚本
⭐ infer.py

import numpy as np
from rknnlite.api import RKNNLite

# 初始化
rknn = RKNNLite()

# 加载模型
ret = rknn.load_rknn('model.rknn')
if ret != 0:
    print('Load failed')
    exit()

# 初始化NPU
ret = rknn.init_runtime()
if ret != 0:
    print('Runtime init failed')
    exit()

# 构造输入（根据你模型）
input_data = np.random.randn(1,1,8,750).astype(np.float32)

# 推理
outputs = rknn.inference(inputs=[input_data])

print(outputs)

rknn.release()

运行：

python infer.py

⑤ 注意事项
报错：

W rknn-toolkit-lite2 version: 2.3.2 Runtime init failed E Catch exception when init runtime! E Traceback (most recent call last): File "/home/cat/anaconda3/lib/python3.8/site-packages/rknnlite/api/rknn_lite.py", line 148, in init_runtime self.rknn_runtime = RKNNRuntime(root_dir=self.root_dir, target=target, device_id=device_id, File "rknnlite/api/rknn_runtime.py", line 363, in rknnlite.api.rknn_runtime.RKNNRuntime.__init__ File "rknnlite/api/rknn_runtime.py", line 607, in rknnlite.api.rknn_runtime.RKNNRuntime._load_library File "rknnlite/api/rknn_runtime.py", line 583, in rknnlite.api.rknn_runtime.RKNNRuntime._get_rknn_api_lib_path Exception: Can not find dynamic library on RK3588! Please download the librknnrt.so from https://github.com/airockchip/rknn-toolkit2/tree/master/rknpu2/runtime/Linux/librknn_api/aarch64 and move it to director

简单来说：

Exception: Can not find dynamic library on RK3588!

板子上缺少 NPU 运行时动态库

librknnrt.so

解决：

从GitHub 下载 ：
进入：rknpu2/runtime/Linux/librknn_api/aarch64/
下载：librknnrt.so（版本与 toolkit 2.3.2 匹配）
拷贝到板子

安装：

# 在板子执行：
sudo cp librknnrt.so /usr/lib/
sudo ldconfig
# 或者
sudo cp librknnrt.so /usr/lib/aarch64-linux-gnu/
sudo ldconfig

再次运行：

rknn.init_runtime()

上述内容相关引用与参考链接

野火嵌入式AI应用开发实战指南：https://doc.embedfire.com/linux/rk356x/Ai/zh/latest/lubancat_ai/env/env.html