1. RKNN量化基础

任务目标

本案例将完成一个最小可复现的RKNN模型转换闭环：把一个训练好的PyTorch分类模型（MobileNetV2，单张图片输入）转换为RKNN格式，并验证转换后的模型能够正确执行推理。通过该流程，理解RKNN-Toolkit2的基本工作阶段，以及量化（INT8）在模型转换中的位置与作用。

环境与依赖

  • RKNN-Toolkit2（版本≥1.6.0），安装参考官方文档：pip install rknn-toolkit2
  • PyTorch（≥1.10），用于导出ONNX模型
  • torchvision（提供MobileNetV2预训练权重）
  • ONNX（≥1.12），用于中间格式检查
  • 运行平台：x86 Linux（模型转换在PC端完成），RKNN模型推理可在PC模拟器或RK3588等NPU板端执行

实现流程与关键代码

整个转换流程分为四个阶段：PyTorch模型导出为ONNX → ONNX模型加载到RKNN-Toolkit2 → 量化与RKNN格式转换 → RKNN模型推理验证。其中量化发生在第三步，RKNN-Toolkit2使用quantize参数控制是否执行INT8量化。本案例使用一张真实图片作为量化校准数据（而非随机噪声），以保证量化后的模型精度损失在可接受范围。

以下代码完整实现了上述流程，可直接保存为convert_mobilenetv2.py运行：

import torch
import torchvision.models as models
import numpy as np
from PIL import Image
from rknn.api import RKNN

# ========== 1. PyTorch模型导出为ONNX ==========
def export_onnx():
    model = models.mobilenet_v2(weights=models.MobileNet_V2_Weights.IMAGENET1K_V1)
    model.eval()
    dummy_input = torch.randn(1, 3, 224, 224)
    torch.onnx.export(model, dummy_input, "mobilenetv2.onnx",
                      input_names=["input"], output_names=["output"],
                      opset_version=12)
    print("ONNX导出完成")

# ========== 2. 准备量化校准数据（真实图片） ==========
def load_calibration_data():
    # 使用一张真实图片作为量化校准输入（实际工程建议50~200张）
    img = Image.open("cat.jpg").resize((224, 224))
    img_data = np.array(img).astype(np.float32) / 255.0
    img_data = np.transpose(img_data, (2, 0, 1))  # HWC -> CHW
    img_data = np.expand_dims(img_data, axis=0)    # 增加batch维度
    return [img_data]

# ========== 3. RKNN转换主流程 ==========
def convert_to_rknn():
    # 初始化RKNN对象
    rknn = RKNN(verbose=True)

    # 配置量化参数（INT8量化）
    rknn.config(mean_values=[[0, 0, 0]], std_values=[[1, 1, 1]],
                target_platform="rk3588", quantize_method="normal")

    # 加载ONNX模型
    ret = rknn.load_onnx(model="mobilenetv2.onnx")
    assert ret == 0, "ONNX加载失败"

    # 执行量化与转换（build时自动完成INT8量化）
    ret = rknn.build(do_quantization=True, dataset="dataset.txt")
    assert ret == 0, "RKNN转换失败"

    # 导出RKNN模型文件
    ret = rknn.export_rknn("mobilenetv2.rknn")
    assert ret == 0, "RKNN导出失败"

    # ========== 4. 推理验证（PC模拟器） ==========
    ret = rknn.init_runtime(target=None)  # None表示使用PC模拟器
    assert ret == 0, "模拟器初始化失败"

    # 加载测试图片并推理
    img = Image.open("cat.jpg").resize((224, 224))
    img_data = np.array(img).astype(np.float32) / 255.0
    img_data = np.transpose(img_data, (2, 0, 1))
    img_data = np.expand_dims(img_data, axis=0)

    outputs = rknn.inference(inputs=[img_data])
    pred_class = np.argmax(outputs[0])
    print(f"预测类别索引: {pred_class}")

    # 释放资源
    rknn.release()

if __name__ == "__main__":
    export_onnx()
    # 创建dataset.txt（每行一个校准图片路径）
    with open("dataset.txt", "w") as f:
        f.write("cat.jpg\n")
    convert_to_rknn()

运行与验证

运行步骤：

  1. 准备一张测试图片cat.jpg（任意内容，建议为常见物体类别），放在与脚本相同目录。
  2. 执行python convert_mobilenetv2.py。
  3. 观察终端输出，确认依次出现“ONNX导出完成”、RKNN build过程日志（含量化信息）、以及最终的预测类别索引。

通过条件：

  • 脚本无报错退出，生成mobilenetv2.onnx和mobilenetv2.rknn两个文件。
  • RKNN推理输出的预测类别索引与PyTorch模型直接推理该图片得到的结果一致（或属于同一语义类别，如都为“猫”）。
  • RKNN build日志中显示量化已完成（出现类似“Quantizing”或“INT8”相关输出），说明量化步骤确实执行。

量化作用说明

在上述流程中，rknn.build(do_quantization=True)将FP32权重与激活值量化为INT8。量化后的RKNN模型体积约为原来的1/4，在NPU上的推理速度显著提升。校准数据（dataset.txt指向的图片）用于统计激活值的动态范围，从而确定量化缩放系数。若省略校准数据或使用随机噪声，模型精度可能大幅下降——这是量化流程中最需要关注的实际问题。

← 上一章 📖 返回目录 下一章 →