440 lines
14 KiB
Markdown
440 lines
14 KiB
Markdown
# 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 频率 |
|