📚 课程目录第1章 RKNN量化基础

1. RKNN量化基础

任务目标

本案例将完成一个最小可复现的RKNN量化闭环：把一个训练好的PyTorch分类模型（ResNet18）转换为RKNN格式，执行INT8量化，并验证量化后模型在模拟器上的推理输出与原始PyTorch模型一致。通过该流程掌握rknn-toolkit2的基本API调用与量化参数含义。

环境依赖

RKNN-Toolkit2（版本≥1.6.0），安装方式：pip install rknn-toolkit2
PyTorch（CPU版本即可）、torchvision、numpy、Pillow
操作系统：Ubuntu 20.04/22.04 x86_64（或Windows WSL）

实现代码

以下代码完成完整转换流程：加载预训练ResNet18 → 导出ONNX → 创建RKNN对象并配置量化 → 加载ONNX并构建RKNN模型 → 导出RKNN文件 → 在模拟器上推理验证。

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

# 1. 导出PyTorch模型为ONNX（使用单张随机输入固定计算图）
model = models.resnet18(pretrained=True)
model.eval()
dummy_input = torch.randn(1, 3, 224, 224)
torch.onnx.export(model, dummy_input, "resnet18.onnx",
                  input_names=["input"], output_names=["output"],
                  opset_version=12)

# 2. 创建RKNN对象并配置量化参数
rknn = RKNN()
rknn.config(mean_values=[[123.675, 116.28, 103.53]],
            std_values=[[58.395, 57.12, 57.375]],
            target_platform="rk3588",
            quantized_dtype="w8a8",
            quantized_algorithm="normal",
            quantized_method="layer")

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

# 4. 构建RKNN模型（执行量化）
ret = rknn.build(do_quantization=True, dataset="dataset.txt")
assert ret == 0, "RKNN构建失败"

# 5. 导出RKNN文件
ret = rknn.export_rknn("resnet18.rknn")
assert ret == 0, "RKNN导出失败"

# 6. 模拟器推理验证
# 准备一张真实图片（此处用随机噪声代替，仅验证流程）
img = np.random.randint(0, 255, (224, 224, 3), dtype=np.uint8)
img = Image.fromarray(img).resize((224, 224))
img = np.array(img).astype(np.float32)

# 初始化模拟器并推理
rknn.init_runtime()
outputs = rknn.inference(inputs=[img])
print("RKNN输出shape:", outputs[0].shape)
print("RKNN输出前5个值:", outputs[0][0][:5])

# 对比PyTorch原始输出
with torch.no_grad():
    torch_out = model(torch.from_numpy(img.transpose(2, 0, 1)).unsqueeze(0).float())
print("PyTorch输出前5个值:", torch_out[0][:5].numpy())

rknn.release()

运行前需在代码同目录创建dataset.txt文件，内容为用于量化的图片路径列表（每行一个路径），例如：

./img1.jpg
./img2.jpg
./img3.jpg

运行方式

将上述代码保存为convert_rknn.py，并准备至少3张真实图片（建议来自ImageNet验证集或任意自然图像）用于量化校准，将路径写入dataset.txt。
执行：python convert_rknn.py

验证方法与判定标准

流程通过：程序无异常退出，终端打印RKNN输出shape: (1, 1000)，且当前目录生成resnet18.rknn文件。
量化正确性：比较RKNN输出与PyTorch输出的前5个值，两者最大类别索引（argmax）应一致。由于INT8量化存在精度损失，允许数值偏差，但分类结果不应改变。
量化参数生效：在rknn.config()中修改quantized_dtype="w8a16"后重新运行，对比两次输出的数值差异——若差异明显，说明量化位宽设置已生效。

关键参数说明

mean_values/std_values：与PyTorch预训练模型的标准化参数一致（ImageNet统计值），必须匹配否则推理结果错误。
quantized_dtype="w8a8"：权重和激活均量化为INT8；改为"w8a16"则激活保留FP16，精度更高但推理速度略慢。
quantized_method="layer"：按层独立计算量化参数；可选"channel"按通道量化，精度更高但模型体积略增。
dataset.txt：提供5～10张代表性图片即可获得稳定量化效果，图片无需标注。

常见问题

ONNX导出失败：确认torchvision版本与PyTorch匹配，或改用opset_version=11重试。
模拟器推理报错：检查target_platform是否与当前RKNN-Toolkit2版本支持的平台一致（如rk3588需1.6.0以上版本）。
输出全为0或NaN：多为mean/std设置错误或输入图片未正确resize到224×224。