本文详细记录一次将 PyTorch 模型 (.pth.tar) 转换为 Rockchip 平台专用 RKNN 模型 (.rknn) 的完整流程。涉及环境配置、依赖冲突以及最终的解决方案，希望能为面临类似任务的开发者提供一份详尽的参考。
0 前言
0.1 最终目标
源模型: 一个基于 PyTorch、以 ResNet 为骨干网络的检测模型。
目标格式: 适用于 Rockchip RK3576 芯片的 RKNN 2.1 版本模型。
核心需求: 模型需经过 INT8 量化以在嵌入式端达到高性能。
0.2 最终成功的环境配置
经过多次尝试，以下是最终确定的、能够稳定运行所有转换任务的环境版本组合。
基础环境:
Linux (e.g., Ubuntu)
Python: 3.6
核心AI框架:
PyTorch: 1.13.1 (cu117)
TorchVision: 0.14.1
TorchAudio: 0.13.1
模型转换工具:
ONNX: 1.12.0
ONNX Runtime: 1.12.0
RKNN-Toolkit2: 2.3.2
关键依赖库:
NumPy: <2.0
OpenCV: 4.8.0.76
Scikit-image, Matplotlib, TQDM, Pandas等
1 完整安装与配置流程
这是一个从零开始的、最可靠的安装步骤，它通过手动控制所有关键软件包的版本来避免自动依赖解析带来的冲突。
1.1 创建并激活 Conda 环境
选择 Python 3.6 以确保最佳的兼容性，创建一个只包含Python 3.6和pip的最小环境
conda create -n rknn_py36_env python=3.6 pip -y
激活新创建的环境
conda activate rknn_py36_env
1.2 安装 RKNN-Toolkit2 及其 cp36 依赖
现在，我们在这个新的 Python 3.6 环境中，严格按照 rknn-toolkit2 为 cp36 提供的文件进行安装。rknn-toolkit2工具下载链接：cknn
进入 rknn-toolkit2 的解压目录：
cd ~/your/rknn-toolkit2-v2.3.2-2025-04-09/
安装 cp36 对应的依赖文件：
pip install -r rknn-toolkit2/packages/x86_64/requirements_cp36-2.3.2.txt
（如果下载慢，可以在命令末尾加上 -i https://pypi.tuna.tsinghua.edu.cn/simple 使用镜像）
安装 cp36 对应的 .whl 主程序文件：
pip install rknn-toolkit2/packages/x86_64/rknn_toolkit2-2.3.2-cp36-cp36m-manylinux_2_17_x86_64.manylinux2014_x86_64.whl
1.3 安装其余需要的库
最后补上脚本本身需要的其他库
pip install opencv-python matplotlib scikit-image
1.4 准备模型转换所需文件
准备校准数据集。
回到您的项目主目录 (请替换为您的实际路径)
cd /path/to/your/main/
准备量化校准集 (假设图片已在 calibration/images/ 目录下)
find calibration/images -type f > calibration.lst
准备转换脚本 convert_to_rknn.py (内容见下一节)
1.5 执行模型转换
由于系统 PATH 变量可能存在问题，建议使用Python解释器的绝对路径来确保调用正确的环境。
查找您的Python解释器绝对路径
which python
(假设输出为 /path/to/your/anaconda3/envs/rknn_py36_env/bin/python)
使用绝对路径运行脚本 (请替换为上一步的实际输出)
/path/to/your/anaconda3/envs/rknn_py36_env/bin/python convert_to_rknn.py
成功后，项目目录下会生成最终的 acm_model_256.rknn 文件。
2 模型转换脚本
点击展开查看完整的Python脚本
import os
import torch
from net import Net
from rknn.api import RKNN
import sys
# ======================= 配置区 (请根据您的实际情况检查) =======================
# 1. PyTorch模型权重文件路径
PYTORCH_MODEL_PATH = 'xxx.pth.tar'
# 2. 中间ONNX模型的文件名
ONNX_MODEL_PATH = 'xxx.onnx'
# 3. 最终RKNN模型的文件名
RKNN_MODEL_PATH = 'xxx.rknn'
# 4. 模型的输入尺寸 (批次, 通道, 高, 宽)
INPUT_SIZE = (1, 1, 256, 256)
# 5. 目标嵌入式平台
TARGET_PLATFORM = 'rk3576'
# 6. 是否进行INT8量化 (强烈建议开启以获得最佳性能)
DO_QUANTIZATION = True
CALIBRATION_DATASET_PATH = './calibration.lst'
# ========================================================================
def export_onnx_model():
"""
第一步：将 PyTorch 模型 (.pth.tar) 转换为 ONNX 格式 (.onnx)
"""
print("--- 步骤 1: 开始导出 ONNX 模型 ---")
try:
model = Net(model_name='ACM', mode='test')
except TypeError:
print("\n[错误] Net 或 model 类初始化失败。")
sys.exit(1)
# 加载模型权重
device = torch.device('cpu')
try:
model.load_state_dict(torch.load(PYTORCH_MODEL_PATH, map_location=device)['state_dict'])
except Exception as e:
print(f"\n[错误] 加载 PyTorch 权重失败: {e}")
print(f"请确认权重文件路径 '{PYTORCH_MODEL_PATH}' 是否正确。")
sys.exit(1)
model.eval()
print("PyTorch 模型加载成功。")
# 创建一个符合模型输入的虚拟张量
dummy_input = torch.randn(INPUT_SIZE, device=device)
# 导出模型为 ONNX 格式
try:
# 使用 opset_version=11 以保证最佳兼容性
torch.onnx.export(model,
dummy_input,
ONNX_MODEL_PATH,
opset_version=11,
input_names=['input'],
output_names=['output'])
except Exception as e:
print(f"\n[错误] 导出 ONNX 失败: {e}")
sys.exit(1)
print(f"ONNX 模型已成功导出到: {ONNX_MODEL_PATH}")
return ONNX_MODEL_PATH
def export_rknn_model(onnx_path):
"""
第二步：将 ONNX 模型 (.onnx) 转换为 RKNN 格式 (.rknn)
"""
print("\n--- 步骤 2: 开始转换 RKNN 模型 ---")
# 初始化 RKNN 对象
rknn = RKNN(verbose=True)
# 配置 RKNN
print("--> 正在配置 RKNN...")
# 对于[0, 255]的单通道输入，归一化到[-1, 1]的mean/std值为127.5
rknn.config(mean_values=[[127.5]],
std_values=[[127.5]],
target_platform=TARGET_PLATFORM,
quantized_dtype='asymmetric_quantized-8')
print("配置完成。")
# 加载 ONNX 模型
print(f"--> 正在加载 ONNX 模型: {onnx_path}")
ret = rknn.load_onnx(model=onnx_path)
if ret != 0:
print("加载 ONNX 模型失败!")
rknn.release()
return
print("加载 ONNX 成功。")
# 构建 RKNN 模型 (包含量化)
print("--> 正在构建 RKNN 模型 (此步骤可能需要几分钟)...")
if DO_QUANTIZATION and not os.path.exists(CALIBRATION_DATASET_PATH):
print(f"\n[错误] 找不到量化校准集文件: {CALIBRATION_DATASET_PATH}")
print("请先创建该文件，或在脚本中设置 DO_QUANTIZATION = False")
rknn.release()
return
ret = rknn.build(do_quantization=DO_QUANTIZATION, dataset=CALIBRATION_DATASET_PATH if DO_QUANTIZATION else None)
if ret != 0:
print("构建 RKNN 模型失败!")
rknn.release()
return
print("构建 RKNN 成功。")
# 导出 RKNN 模型
print(f"--> 正在导出 RKNN 模型到: {RKNN_MODEL_PATH}")
ret = rknn.export_rknn(RKNN_MODEL_PATH)
if ret != 0:
print("导出 RKNN 模型失败!")
rknn.release()
return
print("导出 RKNN 成功！")
# 释放资源
rknn.release()
if __name__ == '__main__':
onnx_file = export_onnx_model()
if onnx_file and os.path.exists(onnx_file):
export_rknn_model(onnx_file)
3 排错历程与问题归纳
| 错误现象 (Symptom) | 根本原因 (Root Cause) | 最终解决方案 (Final Solution) |
|---|---|---|
| ModuleNotFoundError (e.g., torch, cv2) | Conda环境中缺少必需的Python库。 | 在激活的环境中使用 pip install 或 conda install 安装缺失的包。 |
| FileNotFoundError (e.g., config.json, 图片, .lst文件) | 文件路径错误、文件名不匹配、或数据加载代码有缺陷。 | 修正脚本/配置文件中的路径，重命名文件，或修改数据加载器代码以正确拼接扩展名。 |
| IndentationError | 编辑代码时混用了Tab和空格，或意外删除了函数体。 | 使用编辑器的“将缩进转换为空格”功能统一代码格式，或补全缺失的代码。 |
| pip安装成功但python找不到模块 | PATH环境变量配置错误，导致终端优先调用了Conda环境之外的另一个Python解释器。 | [临时方案]使用Python解释器的绝对路径来执行命令。[根本方案]修改~/.bashrc文件，注释掉错误的PATH设置. |
| Solving environment: failed (Conda错误) | 底层网络问题（代理/防火墙）导致Conda无法下载或解析包索引；或极其复杂的依赖冲突。 | [网络]联系管理员配置网络代理。[依赖]彻底重建环境，并采用分步或手动指定所有版本号的方式安装。 |
| AttributeError: 'onnx' has no attribute 'mapping' | RKNN-Toolkit2代码与过新版的onnx库不兼容。 | 将 onnx 和 onnxruntime 降级到 1.12.0 版本。 |
| ...requires onnx>=1.16.1, but you have onnx 1.12.0... | RKNN-Toolkit2安装元数据与其代码的实际运行要求自相矛盾。 | 安装rknn-toolkit2的.whl文件时，使用 --no-dependencies 参数忽略其错误的依赖检查。 |
| RuntimeError: ... no kernel image is available ... | PyTorch版本过旧，不兼容新版GPU的计算能力。 | 安装与GPU硬件匹配的、更新的PyTorch GPU版本。 |
| A module that was compiled using NumPy 1.x cannot be run in NumPy 2.0.x | NumPy 2.0 引入了不向后兼容的API改动。 | 将NumPy降级到2.0以下的版本 (pip install numpy"<2")。 |
4 潜在问题与总结
-
Shell环境 PATH 变量：本次排错最核心的障碍之一。如果服务器的 ~/.bashrc 配置文件没有被修正，未来任何新建的Conda环境都可能遇到 python 命令指向错误位置的问题，需要持续使用绝对路径来规避。
-
服务器网络限制：多次 conda 和 pip 安装失败表明，服务器访问外部网络资源可能受限。配置网络代理是未来最能保证安装顺畅的方法。
-
本次环境配置在python3.9、python3.10、python3.11均未成功，反复出现不匹配错误，使用python3.6后成功，但未究其本质原因。
posted on
浙公网安备 33010602011771号