📚 课程目录第3章 ONNX模型RKNN量化转换

3. ONNX模型RKNN量化转换

任务目标

本案例将完成一个ONNX分类模型的RKNN格式转换与量化全流程：加载ONNX模型，执行INT8量化转换，生成RKNN文件，并验证转换后模型的基本推理可用性。通过本案例，掌握RKNN-Toolkit2中量化转换的核心参数配置与调用方法。

环境依赖

RKNN-Toolkit2（版本≥1.6.0），安装方法参考官方文档
Python 3.8+，依赖库：rknn-toolkit2、onnx、numpy、Pillow
测试平台：PC（x86 Linux/Windows）或RK3588等RKNN目标平台
量化数据集：至少20~50张代表性图片（本案例使用CIFAR-10中抽取的32张图片作为校准集）

关键实现

以下代码完成ONNX模型加载、量化配置、转换与保存的完整流程。核心量化参数包括：mean_values/std_values（输入归一化）、quantized_dtype（INT8量化）、quantized_algorithm（归一化算法）。校准数据集通过Dataset类加载，用于统计激活值范围。

import os
import numpy as np
from PIL import Image
from rknn.api import RKNN

# ========== 配置路径 ==========
ONNX_MODEL = 'resnet18_cifar10.onnx'
RKNN_MODEL = 'resnet18_cifar10_int8.rknn'
DATASET_TXT = 'quant_dataset.txt'   # 每行一张图片路径
IMG_SIZE = (32, 32)                  # CIFAR-10 输入尺寸

# ========== 1. 创建RKNN对象 ==========
rknn = RKNN(verbose=True)

# ========== 2. 配置量化参数 ==========
# mean_values/std_values: CIFAR-10 数据集的归一化参数
# quantized_dtype: 指定INT8量化
# quantized_algorithm: 'normal' 为逐通道归一化量化
ret = rknn.config(
    mean_values=[[0.4914, 0.4822, 0.4465]],
    std_values=[[0.2470, 0.2435, 0.2616]],
    target_platform='rk3588',
    quantized_dtype='w8a8',
    quantized_algorithm='normal',
    quantized_method='layer',
    optimization_level=1
)
assert ret == 0, f'config failed: {ret}'

# ========== 3. 加载ONNX模型 ==========
ret = rknn.load_onnx(model=ONNX_MODEL)
assert ret == 0, f'load_onnx failed: {ret}'

# ========== 4. 构建RKNN模型（执行量化） ==========
# do_quantization=True 时使用校准集进行INT8量化
ret = rknn.build(do_quantization=True, dataset=DATASET_TXT)
assert ret == 0, f'build failed: {ret}'

# ========== 5. 导出RKNN文件 ==========
ret = rknn.export_rknn(RKNN_MODEL)
assert ret == 0, f'export failed: {ret}'

# ========== 6. 初始化运行时并验证推理 ==========
ret = rknn.init_runtime(target=None)  # None表示在PC上模拟运行
assert ret == 0, f'init_runtime failed: {ret}'

# 加载一张测试图片并推理
img = Image.open('test_sample.png').resize(IMG_SIZE)
img = np.array(img).astype(np.uint8)
img = np.expand_dims(img, axis=0)  # 形状: (1, 32, 32, 3)

outputs = rknn.inference(inputs=[img])
print('量化模型推理输出 shape:', outputs[0].shape)
print('Top-1 类别索引:', int(np.argmax(outputs[0][0])))

# ========== 7. 释放资源 ==========
rknn.release()

校准数据集准备

量化校准集用于统计各层激活值的动态范围。创建quant_dataset.txt，每行写入一张图片的路径或相对路径（相对于执行脚本的目录）：

# quant_dataset.txt 示例（前3行）
calib_imgs/img_0001.png
calib_imgs/img_0002.png
calib_imgs/img_0003.png
...

校准图片应从训练集或与真实推理分布相近的图片中随机抽取，建议覆盖各类别且数量不少于20张。图片格式支持JPG/PNG，代码内部会自动缩放到IMG_SIZE。

运行与验证

将ONNX模型、校准图片及quant_dataset.txt置于同一目录，修改代码开头的路径变量。
执行脚本：python convert_rknn.py（假设保存为上述文件名）。
观察终端日志：出现build success和export success即表示转换完成。
验证推理输出：脚本会打印输出张量的shape和Top-1类别索引。若输出shape为(1, 10)（CIFAR-10类别数），且类别索引在0~9范围内，说明RKNN模型可正常推理。

通过判定标准

生成resnet18_cifar10_int8.rknn文件，文件大小明显小于原始ONNX模型（通常为1/3~1/4）。
脚本无报错退出，assert全部通过。
推理输出张量形状与ONNX模型输出一致（CIFAR-10为(1, 10)）。
量化后模型精度损失在可接受范围（可通过与FP32模型在相同测试集上的准确率对比验证，本案例不展开精度评估）。

常见问题

build失败提示dataset路径错误：检查quant_dataset.txt中的图片路径是否相对于当前工作目录正确。
量化后精度显著下降：增加校准集图片数量至100张以上，或检查mean_values/std_values是否与训练时一致。
init_runtime失败：PC模拟运行需安装RKNN-Toolkit2的PC依赖（如onnxruntime）；若在开发板上运行，需将target参数改为对应平台并连接设备。