14 KiB
14 KiB
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 系统要求
# 推荐: 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
# 从 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 安装依赖
pip install onnx onnxruntime numpy opencv-python
# 如果 GPU 可用，安装 onnxruntime-gpu 以加速仿真
三、转换流程详解
3.1 完整转换脚本
code/convert_to_rknn.py：
"""
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 运行转换
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 混合量化示例
# 在 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 运行：
# 在 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 频率 |