RKNN‑Toolkit2 模型转换全流程

RKNN‑Toolkit2 是瑞芯微官方提供的模型转换工具，能够把 PyTorch、ONNX、TensorFlow 等主流框架训练得到的模型，转换为.rknn格式，运行在 RK3588、RK3566、RV1106 等瑞芯微 NPU 硬件上。本文梳理完整转换流程、核心 API、关键参数与完整可运行示例，帮助快速完成边缘 AI 模型转换部署

一、支持的深度学习框架

RKNN‑Toolkit2 支持多种主流框架模型导入，各框架推荐版本如下：

二、RKNN 模型转换整体流程

整体分为 6 大步骤：创建 RKNN 对象 → 配置 RKNN 模型 → 加载原始模型 → 构建 RKNN 模型 → 导出 RKNN 模型 → 释放 RKNN 对象。

创建 RKNN 对象：初始化 RKNN 实例，可以开启日志输出，方便调试排错。

配置 RKNN 模型rknn.config()：设置输入预处理（均值、归一化、RGB/BGR 转换）、量化类型、量化算法、目标芯片平台、模型优化等级等。

加载对应框架模型：调用对应接口load_pytorch/load_onnx/load_tflite等读取训练好的模型文件。

构建 RKNN 模型rknn.build()：执行模型量化、使用校准数据集生成量化参数、完成模型编译构建。

导出 RKNN 模型rknn.export_rknn()：输出.rknn文件，用于开发板端部署。

释放 RKNN 对象rknn.release()：释放内存资源，避免内存泄漏。

三、核心 API 详解

3.1 对象初始化与释放

使用 RKNN 所有接口前，必须初始化RKNN()对象；使用完毕调用release()释放资源。

verbose=True：终端打印详细日志；

verbose_file：将日志保存到本地日志文件，方便排查转换报错。

c 复制代码

from rknn.api import RKNN
# 日志输出到终端，同时写入log.txt
rknn = RKNN(verbose=True, verbose_file="./log.txt")
# 业务逻辑完成之后释放
rknn.release()

3.2 rknn.config () 模型配置（重中之重）

config接口用来做模型预处理、量化策略、目标硬件平台等配置，部分关键参数解析：

mean_values / std_values

输入图像预处理，对应训练时的减均值、除以标准差。格式为嵌套列表，支持多输入模型。

注意：quant_img_RGB2BGR开启时，会先做 RGB2BGR，再执行减均值。

quant_img_RGB2BGR

量化阶段是否做 RGB 转 BGR，常用于 Caffe 模型。

⚠️重要提醒：该参数仅作用于量化校正阶段，不会保存进 rknn 模型文件。部署推理时需要用户自己保证输入图像通道顺序正确。npy 格式量化数据会忽略该配置。

quantized_dtype 量化类型

w8a8（默认）：权重 8bit、激活 8bit 非对称量化，最常用；

w4a16：4bit 权重，16bit 激活；

w8a16：8bit 权重，16bit 激活；

w16a16i、w16a16i_dfp：16bit 定点量化，部分平台支持。

quantized_algorithm 量化算法

normal：速度快，推荐 20‑100 张校准图片；

mmse：精度更高，迭代计算，速度慢，20‑50 张；

kl_divergence：精度与速度折中，适合特征分布不均匀；

gdq：仅 w4a16/w8a16 生效，需要 200 张以上数据集。

quantized_method 量化粒度

layer：每层权重一套量化参数；

channel（默认）：每个输出通道一套量化，精度更高；

group{SIZE}：分组量化，仅 w4a16 可用，进一步提升精度。

target_platform

指定目标芯片，如rk3588、rk3566、rv1106，大小写不敏感。必须和实际部署硬件一致。

optimization_level

模型优化等级，0 关闭全部优化，3 开启全部优化，默认 3。遇到转换后精度异常，可以降低等级尝试。

示例 config 代码

c 复制代码

rknn.config(
    mean_values=[[123.675, 116.28, 103.53]],
    std_values=[[58.395, 58.395, 58.395]],
    quant_img_RGB2BGR=False,
    quantized_dtype="w8a8",
    quantized_algorithm="normal",
    quantized_method="channel",
    target_platform="rk3588",
    optimization_level=3
)

3.3 模型加载接口

不同框架调用不同 load 接口：

PyTorch：rknn.load_pytorch(model="xxx.pt", input_size_list= $\[1,3,224,224$ ])，模型需要导出为 TorchScript；

ONNX：rknn.load_onnx(model="xxx.onnx")，工程最常用；

TensorFlow：rknn.load_tensorflow()；

TFLite：rknn.load_tflite()；

Caffe：rknn.load_caffe()；

Darknet：rknn.load_darknet()

c 复制代码

ret = rknn.load_pytorch(
    model="./resnet18.pt",
    input_size_list=[[1, 3, 224, 224]]
)
if ret !=0:
    print("模型加载失败")
    exit(-1)

3.4 rknn.build () 构建 RKNN 模型

build 阶段完成模型编译、量化。关键参数：

do_quantization=True：开启 INT8 量化，关闭则为 FP16 浮点模型；

dataset：量化校正数据集 txt 文件路径。txt 每行写一张图片 /npy 的路径；多输入模型一行写多个文件路径，空格分隔。

提示：校正数据集尽量选用和真实业务场景一致的数据，一般 20‑100 张，保证量化精度；

rknn_batch_size：设置推理 batch，多核平台可以提升硬件利用率，会改变模型输入输出 shape；

auto_hybrid：自动混合量化，部分 OP 自动切换 FP16，解决量化溢出、精度掉点。

c 复制代码

ret = rknn.build(
    do_quantization=True,
    dataset="./dataset.txt"
)
if ret !=0:
    print("build失败")
    exit(-1)

3.5 rknn.export_rknn () 导出模型

将构建完成的模型保存为.rknn文件，给到板卡部署使用。

c 复制代码

ret = rknn.export_rknn(export_path="./resnet18.rknn")

四、完整转换示例（PyTorch → RKNN）

c 复制代码

from rknn.api import RKNN

if __name__ == '__main__':
    # 1.初始化对象
    rknn = RKNN(verbose=True, verbose_file="convert.log")

    # 2.模型配置
    rknn.config(
        mean_values=[[123.675, 116.28, 103.53]],
        std_values=[[58.395, 58.395, 58.395]],
        quantized_dtype="w8a8",
        quantized_algorithm='normal',
        quantized_method='channel',
        quant_img_RGB2BGR=False,
        target_platform="rk3588",
        float_dtype="float16",
        optimization_level=3,
        compress_weight=False
    )

    # 3.加载pytorch torchscript模型
    ret = rknn.load_pytorch(
        model="./resnet18.pt",
        input_size_list=[[1, 3, 224, 224]]
    )
    if ret != 0:
        print("加载模型失败！")
        rknn.release()
        exit(-1)

    # 4.构建模型，开启量化
    ret = rknn.build(do_quantization=True, dataset="dataset.txt")
    if ret != 0:
        print("build构建失败！")
        rknn.release()
        exit(-1)

    # 5.导出rknn模型
    ret = rknn.export_rknn(export_path="resnet18.rknn")
    if ret ==0:
        print("✅ RKNN模型导出成功 resnet18.rknn")

    # 6.释放资源
    rknn.release()

五、踩坑与注意事项

预处理参数必须和训练完全对齐：mean_values、std_values、RGB/BGR 通道顺序错误，会直接导致推理结果完全错误。

quant_img_RGB2BGR仅用于量化阶段，不会写入 rknn 模型，板卡部署时图像通道顺序需要自己处理。

量化数据集尽量使用真实业务图片，不要用随机图片，否则量化后精度严重下降。

target_platform一定要写对，不同芯片的 rknn 模型不能互相通用。

PyTorch 模型不能直接传入.pth权重文件，必须导出为 TorchScript 格式才能被load_pytorch加载。

转换结束务必调用release()释放 RKNN 对象，防止内存占用过高。

量化后精度掉点，可以尝试更换quantized_algorithm算法、开启auto_hybrid自动混合量化、增加校准图片数量。

六、总结

RKNN‑Toolkit2 模型转换流程可以概括为：初始化对象 → config 配置预处理与量化参数 → load 加载模型 → build 编译量化 → export 导出 rknn 文件 → release 释放资源。

其中最容易出问题的地方集中在预处理参数、量化数据集、目标平台设置。理解各个 config 参数含义，配合日志输出调试，就可以顺利完成深度学习模型向瑞芯微 NPU 的迁移。