# 03 — RKNN 模型转换与量化详解 ## 一、RKNN-Toolkit2 概述 ``` PC 端 (x86_64 Linux) RK3568 板端 ┌──────────────────────┐ ┌───────────────────┐ │ RKNN-Toolkit2 │ │ librknnrt.so │ │ │ │ │ │ PyTorch/ONNX/TF │ │ .rknn 文件 │ │ │ │ │ │ │ │ ▼ │ │ ▼ │ │ 图优化 + 量化 + 编译 │ ─── .rknn ───► │ RKNN Runtime │ │ │ (模型文件) │ │ │ │ 仿真运行 (可选) │ │ ▼ │ │ │ │ NPU 硬件执行 │ └──────────────────────┘ └───────────────────┘ ``` **关键理解**: - **RKNN-Toolkit2** 只在 PC 端用一次(模型转换),不跑在板子上 - 板端只需要 **librknnrt.so** + 转换好的 **.rknn 文件** - 转换过程:ONNX → 图优化 → 量化校准 → 编译 → .rknn ## 二、环境安装 ### 2.1 系统要求 ```bash # 推荐: Ubuntu 20.04/22.04 x86_64 # Python: 3.8 (推荐) 或 3.9 # 内存: 8GB+ (量化校准需要加载数据集) # 创建虚拟环境 python3.8 -m venv rknn_env source rknn_env/bin/activate ``` ### 2.2 安装 RKNN-Toolkit2 ```bash # 从 Rockchip SDK 获取 whl 包 # SDK 路径: external/rknn-toolkit2/rknn-toolkit2/packages/ # 安装(选择对应 Python 版本的包) pip install rknn_toolkit2-1.5.2-cp38-cp38-linux_x86_64.whl # 验证安装 python -c "from rknn.api import RKNN; print('OK')" # 预期输出: OK ``` ### 2.3 安装依赖 ```bash pip install onnx onnxruntime numpy opencv-python # 如果 GPU 可用,安装 onnxruntime-gpu 以加速仿真 ``` ## 三、转换流程详解 ### 3.1 完整转换脚本 `code/convert_to_rknn.py`: ```python """ ONNX → RKNN 模型转换 + INT8 量化 用于电力波形故障识别模型的转换 转换链路: PyTorch model → ONNX → (本脚本) → RKNN 关键参数说明: - do_quantization=True: 开启 INT8 量化 - dataset: 校准数据集 (100~500 张典型样本) - target_platform: 'rk3562' 或 'rk3568' (RKNN Toolkit 2.1.0+ 支持) """ import numpy as np import os from rknn.api import RKNN # ============================================================ # 配置区 (请根据实际情况修改) # ============================================================ # 模型路径 ONNX_MODEL = 'checkpoints/fault_1dcnn.onnx' RKNN_MODEL = 'deploy/fault_1dcnn.rknn' # 模型参数 (必须与训练时一致) INPUT_CH = 1 SEQ_LEN = 128 NUM_CLASSES = 5 # 归一化参数 (从训练集统计) # 注意: 这些参数在模型外部用 C 代码计算,不放在 ONNX 里 MEAN = 0.0 # 电流均值 (μ),如果做 Z-score 归一化 STD = 1.0 # 电流标准差 (σ) # 如果训练时已归一化到 [0,1] 或 [-1,1],则 mean=0, std=1 # 校准数据集配置 CALIB_DATASET = 'calib_data/' # 校准样本目录 CALIB_FILE = 'calib_dataset.txt' # 校准文件列表 (每行一个 .npy 路径) TARGET = 'rk3562' # 目标平台 (RK3568 也填 rk3562,NPU IP 相同) QUANT_DTYPE = 'asymmetric_quantized-u8' # INT8 非对称量化 # ============================================================ # 1. 准备校准数据集 # ============================================================ def prepare_calibration_data(): """ 从训练集中提取 100~500 个样本作为校准数据。 校准样本应覆盖所有故障类型和正常运行状态。 每个样本保存为独立的 .npy 文件,形状与模型输入一致。 """ os.makedirs(CALIB_DATASET, exist_ok=True) # 如果有现成的校准数据集,跳过 if os.path.exists(CALIB_FILE): with open(CALIB_FILE) as f: lines = f.readlines() print(f"校准数据集已就绪: {len(lines)} 个样本") return # 否则从训练数据提取 # 这里假设有 train_data.npy 和 train_labels.npy print("正在从训练集提取校准样本...") try: data = np.load('data/train_data.npy').astype(np.float32) labels = np.load('data/train_labels.npy') # 确保是 (N, 1, 128) 格式 if data.ndim == 2: data = data[:, np.newaxis, :] # 每类均匀采样 indices = [] for cls in range(NUM_CLASSES): cls_idx = np.where(labels == cls)[0] n_sample = min(100, len(cls_idx)) # 每类 100 个 sampled = np.random.choice(cls_idx, n_sample, replace=False) indices.extend(sampled) np.random.shuffle(indices) # 限制总数 500 个以内 indices = indices[:500] with open(CALIB_FILE, 'w') as f: for i, idx in enumerate(indices): fname = f'{CALIB_DATASET}/calib_{i:04d}.npy' np.save(fname, data[idx]) f.write(f'{fname}\n') print(f"校准数据集已生成: {len(indices)} 个样本") except FileNotFoundError: print("WARNING: 未找到训练数据,生成随机校准样本(仅用于流程验证)") # 生成随机校准样本(实际部署时请替换为真实数据) with open(CALIB_FILE, 'w') as f: for i in range(200): fname = f'{CALIB_DATASET}/calib_{i:04d}.npy' dummy = np.random.randn(1, INPUT_CH, SEQ_LEN).astype(np.float32) np.save(fname, dummy) f.write(f'{fname}\n') print("随机校准样本已生成: 200 个") # ============================================================ # 2. 模型转换主流程 # ============================================================ def convert_model(): rknn = RKNN(verbose=True) # ---- 2.1 配置 ---- print("\n[1/5] 配置 RKNN...") ret = rknn.config( mean_values=[[MEAN]], # 每通道的均值 std_values=[[STD]], # 每通道的标准差 target_platform=TARGET, quantized_dtype=QUANT_DTYPE, # 优化选项 optimization_level=3, # 0~3, 3=最激进的图优化 # 内存优化 # single_core_mode=False, # 多核并行, 加速推理 ) if ret != 0: raise RuntimeError(f"rknn.config 失败: {ret}") # ---- 2.2 加载 ONNX ---- print("\n[2/5] 加载 ONNX 模型...") ret = rknn.load_onnx( model=ONNX_MODEL, inputs=['waveform'], input_size_list=[[INPUT_CH, SEQ_LEN]], input_initial_val=[np.random.randn(1, INPUT_CH, SEQ_LEN).astype(np.float32)], ) if ret != 0: raise RuntimeError(f"ONNX 加载失败: {ret}。" "请检查 ONNX 模型的 opset 版本和算子兼容性。") # ---- 2.3 构建模型 (含量化) ---- print("\n[3/5] 构建模型 (INT8 量化)...") ret = rknn.build( do_quantization=True, dataset=CALIB_FILE, # 量化精度调优选项 quantized_method='layer', # 'layer' 逐层量化 (推荐) # quantized_method='channel', # 'channel' 逐通道量化 (精度更高但耗时) ) if ret != 0: raise RuntimeError(f"模型构建失败: {ret}。" "常见原因: 1) 校准数据量不足 2) 有不支持的算子 3) 量化后精度损失过大") # ---- 2.4 导出 RKNN ---- print("\n[4/5] 导出 RKNN 模型...") ret = rknn.export_rknn(RKNN_MODEL) if ret != 0: raise RuntimeError(f"RKNN 导出失败: {ret}") model_size = os.path.getsize(RKNN_MODEL) print(f"RKNN 模型大小: {model_size / 1024:.1f} KB") # ---- 2.5 精度验证 (PC 端仿真) ---- print("\n[5/5] PC 端精度验证...") rknn.init_runtime(target=TARGET) # 加载测试集 try: test_data = np.load('data/test_data.npy').astype(np.float32) test_labels = np.load('data/test_labels.npy') if test_data.ndim == 2: test_data = test_data[:, np.newaxis, :] except FileNotFoundError: print("WARNING: 测试集不存在,用随机数据验证推理流程") test_data = np.random.randn(100, INPUT_CH, SEQ_LEN).astype(np.float32) test_labels = np.random.randint(0, NUM_CLASSES, 100) correct = 0 all_preds = [] MAX_TEST = min(200, len(test_data)) # 测试 200 条以内 for i in range(MAX_TEST): sample = test_data[i:i+1] # (1, C, L) output = rknn.inference(inputs=[sample])[0] # (1, NUM_CLASSES) pred = np.argmax(output) all_preds.append(pred) if pred == test_labels[i]: correct += 1 acc = correct / MAX_TEST print(f"PC 仿真准确率: {acc:.2%} ({correct}/{MAX_TEST})") # 精度下降检查 # 如果仿真准确率比训练时下降超过 2%,需要排查量化问题 if acc < 0.90: print("⚠️ WARNING: 仿真准确率偏低!") print(" 建议: 1) 增大校准数据集 2) 尝试混合量化 3) 考虑 QAT") else: print("✓ 精度验证通过") rknn.release() return acc # ============================================================ # 3. 性能基准测试 (可选) # ============================================================ def benchmark(): """ 在 PC 端模拟 NPU 推理性能(仅供参考,实际板端性能可能不同) """ print("\n===== 性能基准测试 (PC 仿真) =====") rknn = RKNN() rknn.load_rknn(RKNN_MODEL) rknn.init_runtime(target=TARGET, perf_debug=True) dummy = np.random.randn(1, INPUT_CH, SEQ_LEN).astype(np.float32) import time # 预热 for _ in range(10): rknn.inference(inputs=[dummy]) # 计时 times = [] for _ in range(100): t0 = time.perf_counter() rknn.inference(inputs=[dummy]) t1 = time.perf_counter() times.append((t1 - t0) * 1000) print(f"平均推理耗时: {np.mean(times):.2f} ms") print(f"最小推理耗时: {np.min(times):.2f} ms") print(f"最大推理耗时: {np.max(times):.2f} ms") print(f"抖动 (std): {np.std(times):.2f} ms") print(f"fps (等效): {1000 / np.mean(times):.0f} fps") # 打印模型信息 rknn.eval_perf(is_print=True) # 显示各层在 NPU 还是 CPU 上执行 rknn.release() # ============================================================ # 4. 主流程 # ============================================================ if __name__ == '__main__': os.makedirs('deploy', exist_ok=True) # Step 1: 准备校准数据 prepare_calibration_data() # Step 2: 转换模型 acc = convert_model() # Step 3: 性能基准 benchmark() print(f"\n===== 转换完成 =====") print(f"RKNN 模型: {RKNN_MODEL}") print(f"PC 仿真准确率: {acc:.2%}") print(f"下一步: 将 {RKNN_MODEL} 部署到 RK3568 板端") ``` ### 3.2 运行转换 ```bash python convert_to_rknn.py # 输出示例: # [1/5] 配置 RKNN... # [2/5] 加载 ONNX 模型... # I NPUT: ['waveform'] # O UTPUT: ['logits'] # [3/5] 构建模型 (INT8 量化)... # I Build model... # I Quantization... (使用 200 个校准样本) # I Build done. # [4/5] 导出 RKNN 模型... # RKNN 模型大小: 28.5 KB # [5/5] PC 端精度验证... # PC 仿真准确率: 98.50% (197/200) # ✓ 精度验证通过 ``` --- ## 四、量化调优策略 ### 4.1 精度不够怎么办? ``` 精度损失 > 2% 时的排查顺序: 1. 增大校准数据集 └── 从 100 个 → 500 个样本, 确保每类样本均匀分布 2. 改用逐通道量化 └── quantized_method='channel' (精度更高, 模型略大) 3. 混合量化 └── 只对精度敏感层用 FP16, 其他层 INT8 ``` rknn.build(do_quantization=True, ...) # 然后: rknn.hybrid_quantization_step1(dataset) rknn.hybrid_quantization_step2(exported_rknn, target_platform) ``` 4. 量化感知训练 (QAT) └── 最彻底的方式, 在训练时就考虑量化误差 └── PyTorch 中插入 FakeQuantize 节点 → 重新训练 → 导出 └── RKNN-Toolkit2 的 hybrid_quantization 可加载 QAT 模型 ``` ### 4.2 混合量化示例 ```python # 在 build 之后,如果精度不理想: rknn.build(do_quantization=True, dataset=CALIB_FILE) # 分析每层的量化误差 rknn.accuracy_analysis( inputs=['calib_data/calib_0000.npy'], target=TARGET, device_id=None ) # 对误差大的层改为 FP16 # 例如 Conv1 和最后的 FC 层保持 FP16,中间层 INT8 custom_quant = { 'Conv_0': 'fp16', # 输入层保持高精度 'Gemm_0': 'fp16', # 分类头保持高精度 # 其他层默认 INT8 } rknn.hybrid_quantization_step1(dataset, custom_quant) # ... ``` ### 4.3 算子回退处理 如果某个算子在 NPU 上不支持或精度损失太大,可以让它回退到 CPU 运行: ```python # 在 rknn.config() 中指定自定义算子回退 ret = rknn.config( ... custom_string="""custom_op_name:CPU """, # 例如: # custom_string=""" # /block1/Conv_0:CPU # """ ) ``` > 注意:CPU 回退会显著降低推理速度,仅在万不得已时使用。 --- ## 五、常见问题 | 问题 | 原因 | 解决 | |------|------|------| | `load_onnx` 失败 | ONNX opset 不兼容 | 降低 opset 到 11 或 12 | | `build` 内存不足 | 校准数据集太大 | 限制校准样本 500 个以内 | | 仿真精度远低于原模型 | 量化损失 | 使用混合量化或 QAT | | 某些层回退到 CPU | 算子不支持 | 修改原模型的网络结构 | | 推理多次后变慢 | NPU 温度升高, 降频 | 加散热片, 检查 NPU 频率 |