📚 课程目录第4章 RKNN量化配置

4. RKNN量化配置

任务目标

在RKNN模型转换过程中，量化算法的选择直接影响模型精度与推理性能。本案例以YOLOv8s检测模型为例，演示如何在RKNN-Toolkit2中分别配置normal（普通量化）与mmse（最小均方误差量化）两种算法，并对比二者在COCO验证集上的mAP差异。

环境依赖

RKNN-Toolkit2（版本≥1.6.0）
rknn-toolkit2 Python包及其依赖（numpy、onnx、opencv-python等）
YOLOv8s ONNX模型（可从Ultralytics官方导出）
COCO val2017子集（至少100张用于快速验证）

关键完整代码

import os
import cv2
import numpy as np
from rknn.api import RKNN

# ========== 配置区 ==========
ONNX_MODEL = "yolov8s.onnx"
RKNN_MODEL_PREFIX = "yolov8s_quant"
DATASET_TXT = "coco_subset.txt"  # 每行一张图片路径
IMG_SIZE = 640
QUANT_ALGORITHMS = ["normal", "mmse"]  # 待对比的量化算法

def create_dataset_txt(img_dir, output_txt, num_samples=100):
    """从图片目录生成量化校准数据集列表"""
    imgs = sorted(os.listdir(img_dir))[:num_samples]
    with open(output_txt, "w") as f:
        for name in imgs:
            f.write(os.path.join(img_dir, name) + "\n")

def build_and_export(algorithm):
    """使用指定量化算法构建RKNN模型并导出"""
    rknn = RKNN(verbose=False)
    
    # 配置量化参数（核心差异点）
    rknn.config(
        mean_values=[[0, 0, 0]],
        std_values=[[255, 255, 255]],
        target_platform="rk3588",
        quantized_dtype="asymmetric_quantized-8",
        quantized_algorithm=algorithm,   # 关键：选择量化算法
        quantized_method="layer",        # 逐层量化
    )
    
    # 加载ONNX模型
    assert rknn.load_onnx(model=ONNX_MODEL) == 0, "ONNX加载失败"
    
    # 构建（执行量化）
    ret = rknn.build(do_quantization=True, dataset=DATASET_TXT)
    assert ret == 0, f"{algorithm} 量化构建失败"
    
    # 导出RKNN模型
    output_path = f"{RKNN_MODEL_PREFIX}_{algorithm}.rknn"
    assert rknn.export_rknn(output_path) == 0, "导出失败"
    rknn.release()
    return output_path

def evaluate_mAP(rknn_model, val_img_dir, val_label_dir, img_size=640):
    """在COCO子集上评估mAP@0.5:0.95（简化版，仅演示流程）"""
    from pycocotools.coco import COCO
    from pycocotools.cocoeval import COCOeval
    
    rknn = RKNN(verbose=False)
    rknn.load_rknn(rknn_model)
    assert rknn.init_runtime() == 0, "运行时初始化失败"
    
    # 加载COCO标注（需提前准备子集标注文件）
    coco_gt = COCO(val_label_dir)
    detections = []
    
    for img_id in coco_gt.getImgIds():
        img_info = coco_gt.loadImgs(img_id)[0]
        img_path = os.path.join(val_img_dir, img_info["file_name"])
        img = cv2.imread(img_path)
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        img = cv2.resize(img, (img_size, img_size))
        
        # 前处理 + 推理
        outputs = rknn.inference(inputs=[img])
        # 此处需接入YOLOv8后处理（NMS、坐标解码等）
        # 将检测结果转换为COCO格式并加入detections列表
        # ...（后处理代码略，需根据实际输出格式实现）
    
    # COCO评估
    coco_dt = coco_gt.loadRes(detections)
    coco_eval = COCOeval(coco_gt, coco_dt, "bbox")
    coco_eval.evaluate()
    coco_eval.accumulate()
    coco_eval.summarize()
    return coco_eval.stats[0]  # mAP@0.5:0.95

if __name__ == "__main__":
    # 1. 准备校准数据集（若已存在则跳过）
    if not os.path.exists(DATASET_TXT):
        create_dataset_txt("val2017_sample", DATASET_TXT)
    
    # 2. 分别构建两种量化模型
    models = {}
    for algo in QUANT_ALGORITHMS:
        print(f"正在构建 {algo} 量化模型...")
        models[algo] = build_and_export(algo)
    
    # 3. 对比评估（需准备COCO子集标注）
    print("\n===== 量化算法对比结果 =====")
    results = {}
    for algo, model_path in models.items():
        mAP = evaluate_mAP(model_path, "val2017_sample", "annotations/instances_val2017_subset.json")
        results[algo] = mAP
        print(f"{algo}: mAP@0.5:0.95 = {mAP:.4f}")
    
    # 4. 输出对比结论
    diff = results["mmse"] - results["normal"]
    print(f"\nmmse 相对 normal 的mAP差异: {diff:+.4f}")
    print("结论: mmse算法通常能提供更高精度，但转换时间更长。")

运行与验证

准备数据：将COCO val2017中至少100张图片放入val2017_sample/目录，并准备对应的子集标注JSON文件。
执行脚本：python quant_compare.py，脚本将依次完成normal和mmse两种算法的模型构建。
观察输出：终端会打印两种算法的mAP值及差异。若mmse的mAP高于normal，说明mmse量化在本模型上有效。

通过条件：脚本无报错运行完毕，生成yolov8s_quant_normal.rknn和yolov8s_quant_mmse.rknn两个模型文件，且对比结果中mmse的mAP不低于normal（通常高0.5~2个百分点）。

常见问题

校准数据集过少：若mAP波动大，将num_samples增至500~1000张。
后处理未实现：本案例代码中YOLOv8后处理部分需根据实际输出张量形状补充NMS与坐标解码，否则评估函数无法完成。
mmse构建缓慢：mmse算法需要迭代优化，转换时间可能是normal的3~5倍，属正常现象。