背景

RK3588 内置 6 TOPS NPU，支持 int4/int8/int16/FP16 量化推理。但官方工具链文档碎片化严重，从训练到部署每一步都有坑。这篇记录我的完整流程。

环境准备

# 确认 NPU 驱动
ls /dev/dri/renderD128
cat /sys/kernel/debug/rknpu/version

# 安装 rknn-toolkit2
pip install rknn-toolkit2-2.3.2-cp310-cp310-manylinux_2_17_aarch64.whl

# 验证
python3 -c "from rknn.api import RKNN; print('NPU OK')"

⚠️ rknn-toolkit2 2.x 要求 numpy < 2.0，装完记得检查版本。

模型转换

1. 导出 PyTorch 模型为 ONNX

import torch
import torchvision

# 加载模型
model = torchvision.models.resnet50(pretrained=True)
model.eval()

# 导出
dummy_input = torch.randn(1, 3, 224, 224)
torch.onnx.export(
    model, dummy_input, "resnet50.onnx",
    input_names=["input"],
    output_names=["output"],
    dynamic_axes={"input": {0: "batch"}, "output": {0: "batch"}}
)

2. ONNX → RKNN

from rknn.api import RKNN

rknn = RKNN()
rknn.config(mean_values=[123.675, 116.28, 103.53],
            std_values=[58.395, 57.12, 57.375],
            target_platform="rk3588")

rknn.build(do_quantization=True, quantize_dtype="int8")
rknn.export_rknn("resnet50.rknn")

推理验证

from rknn.api import RKNN

rknn = RKNN()
rknn.load_rknn("resnet50.rknn")

ret = rknn.init_runtime()
if ret != 0:
    print("Runtime init failed")
    exit()

# 推理
import numpy as np
input_data = np.random.randn(1, 3, 224, 224).astype(np.float32)
outputs = rknn.inference([input_data])
print("Output shape:", outputs[0].shape)

性能测试

import time
import numpy as np

# Warm up
for _ in range(10):
    _ = rknn.inference([input_data])

# Benchmark
times = []
for _ in range(100):
    start = time.time()
    _ = rknn.inference([input_data])
    times.append(time.time() - start)

print(f"Avg: {np.mean(times)*1000:.2f}ms, FPS: {1/np.mean(times):.1f}")

实测 ResNet50 int8 量化后：单帧推理 8ms，FPS 125，CPU 模式需要 45ms。

常见踩坑

  1. 输入维度不匹配 — rknn.config 的 mean/std 和训练时要一致
  2. 动态 shape — 如果导出时用了 dynamic_axes，转换会变慢，建议固定 batch=1 先跑通
  3. quantize_dtype — int8 精度损失约 1-2%，但速度最快；int16 精度更高但速度慢
  4. NPU 内存不足 — 大模型需要分片或者降低量化精度

后续优化方向

  • 换成 int4 量化，内存占用减半
  • 多线程并发推理
  • 视频流批处理