瑞芯微RV1126平台人脸识别模型转换实战:从ONNX到RKNN的完整流程

最近在给一个嵌入式门禁项目做算法选型,最终敲定了瑞芯微的RV1126。这颗芯片在边缘端AI计算上的性价比确实不错,但真要把训练好的PyTorch或TensorFlow模型部署上去,中间那道“模型转换”的坎儿,着实让我和团队折腾了好一阵子。尤其是人脸识别这类对精度和延迟都极其敏感的应用,从ONNX到RKNN的转换,远不是点几下鼠标就能搞定的事。这里面涉及到量化策略的选择、预处理对齐、运行时内存优化等一系列细节,任何一个环节出岔子,轻则精度暴跌,重则模型根本跑不起来。

今天,我就把自己在RV1126上部署一个人脸识别模型(以典型的ArcFace或类似结构为例)的完整过程、踩过的坑以及验证有效的优化技巧,系统地梳理出来。这份指南面向的是已经具备一定深度学习基础,正在或即将面临嵌入式AI落地挑战的工程师。我们不谈空洞的理论,只聚焦于从拿到ONNX模型开始,到在RV1126上获得稳定可靠推理结果的每一步实操。

1. 转换前的核心准备:理解RV1126与RKNN的约束

在动手写任何转换代码之前,我们必须先搞清楚目标平台的能力边界和规则。RV1126搭载的NPU(神经网络处理单元)有其特定的硬件架构和指令集,这直接决定了RKNN模型格式的种种限制。盲目转换,大概率会事倍功半。

1.1 RV1126 NPU架构与模型兼容性要点

RV1126的NPU算力标称在1.2 TOPS左右,支持INT8/INT16混合量化,这对降低功耗和提升推理速度至关重要。但它并非一个“万能推理引擎”,对模型结构有明确偏好。

  • 算子支持度:这是第一道关卡。并非所有ONNX算子都能被RKNN-Toolkit直接转换。例如,早期版本的RKNN对GridSampleInstanceNorm等算子的支持并不完善。在转换前,务必查阅官方提供的《RKNN-Toolkit OP支持列表》,确认你的模型中所有算子都在支持范围内。一个实用的技巧是,在导出ONNX模型时,尽量使用更通用、更基础的算子组合来替代那些可能不被支持的复杂算子。
  • 输入输出格式:NPU对输入张量的布局(Layout)非常敏感。最常见的图像输入格式是NCHW(批大小、通道、高度、宽度)。如果你的训练模型使用的是NHWC,必须在转换前或转换配置中明确指出并进行转换,否则会导致结果完全错误。
  • 动态形状限制:RV1126的NPU对动态Batch Size或动态分辨率支持有限。在大多数生产场景下,固定输入尺寸是更稳妥的选择。这意味着,在模型训练和导出时,最好就确定好推理时图像的固定尺寸(例如112x112160x160)。

为了更直观地了解常见人脸识别模型在转换时可能遇到的算子兼容性问题,我整理了一个简表:

模型常见层/算子 RKNN支持情况 转换注意事项与替代方案
卷积 (Conv) 完全支持 注意分组卷积(Group Conv)的组数需为NPU友好值(如1, 2, 4, 8, 16)。
批归一化 (BatchNorm) 支持(通常融合) 在转换时,BN层参数通常会与前一层的卷积或全连接层进行融合,以提升效率。确保训练时BN的track_running_stats为True。
激活函数 (ReLU, PReLU等) 广泛支持 PReLU的斜率参数需要正确导出。
全连接层 (Linear/Gemm) 支持 注意大尺寸全连接层可能占用大量内存,需评估RV1126的SRAM大小。
池化层 (MaxPool, AvgPool) 支持 支持常规核尺寸与步长。
Flatten/Reshape 支持 确保形状变换后的张量在内存上是连续的,避免出现非标准布局。
注意力机制模块 视具体算子而定 如SE模块中的全局池化和缩放操作通常支持,但自定义的复杂注意力可能需要拆解为基本算子。

提示:在将PyTorch模型转换为ONNX时,使用torch.onnx.exportopset_version参数不宜过高,选择稳定版本(如opset 11或12)能提高与RKNN-Toolkit的兼容性。

1.2 环境搭建:不仅仅是运行Docker

原始资料里提到了用Docker启动RKNN-Toolkit,这确实是最简单无污染的方式。但如果你想进行更深入的调试,或者需要与自己的开发环境集成,了解其内部构成更有帮助。

官方提供的Docker镜像已经包含了RKNN-Toolkit及其所有依赖。启动命令的核心是挂载USB设备,以便后续连接真实的RV1126开发板进行模型性能分析和调试。

# 这是最基础的启动命令,映射了USB设备并进入交互式终端
docker run -t -i --privileged -v /dev/bus/usb:/dev/bus/usb rockchip-linux/rknn-toolkit:1.7.1 /bin/bash

然而,在实际项目中,我建议对这个流程做两点增强:

  1. 持久化工作目录:添加-v $(pwd)/workspace:/workspace参数,将宿主机的一个目录挂载到容器内。这样,所有生成的模型文件、测试数据和脚本都不会随着容器的退出而消失。
  2. 使用脚本管理:不要每次都输入一长串命令。创建一个简单的start_rknn_docker.sh脚本,将完整的docker run命令写进去,并附带一些初始化操作(如安装额外的Python包)。
#!/bin/bash
# start_rknn_docker.sh
HOST_WORKSPACE="/path/to/your/project"
docker run -t -i --rm \
  --privileged \
  -v /dev/bus/usb:/dev/bus/usb \
  -v $HOST_WORKSPACE:/workspace \
  --name rknn_dev \
  rockchip-linux/rknn-toolkit:1.7.1 \
  /bin/bash -c "cd /workspace && /bin/bash"

这样,每次启动都是一个干净且工作目录固定的环境。

2. 模型转换的实战步骤:从ONNX到RKNN

环境就绪后,就进入了核心的模型转换阶段。这个过程可以分解为加载、配置、构建(量化)和导出四个关键步骤。我将结合一个具体的人脸识别模型face_recog.onnx来详细说明。

2.1 模型探查与预处理对齐

在转换之前,必须清楚你的模型输入输出细节。使用Netron(https://netron.app)打开ONNX模型是一个好习惯。你需要准确记录:

  • 输入节点名:通常是inputdata
  • 输入形状:例如[1, 3, 112, 112],分别代表批大小、通道数、高、宽。
  • 输入数据类型:通常是float32
  • 输出节点名:可能不止一个,对于人脸识别,我们通常关注特征向量输出。

但更重要的是预处理对齐。这是精度损失的主要来源之一。你的训练代码和推理代码必须使用完全相同的图像预处理流程。假设训练时预处理代码如下:

# 训练时的预处理 (PyTorch示例)
def preprocess_train(image):
    # 1. 缩放到固定尺寸,如 112x112
    # 2. 像素值从 [0, 255] 归一化到 [0, 1]
    image = image / 255.0
    # 3. 应用标准化: (x - mean) / std
    # 假设 mean = [0.5, 0.5, 0.5], std = [0.5, 0.5, 0.5]
    mean = torch.tensor([0.5, 0.5, 0.5])
    std = torch.tensor([0.5, 0.5, 0.5])
    image = (image - mean) / std
    return image

那么,在RKNN转换配置中,mean_valuesstd_values参数就必须与此对应。注意,RKNN配置中的meanstd是应用于归一化到[0,1]之后的像素值,并且通常以0-255范围的整数形式表示。计算关系为:rknn_mean = mean * 255, rknn_std = std * 255

根据上面的例子:mean=[0.5,0.5,0.5], std=[0.5,0.5,0.5],那么:

  • rknn_mean = [0.5*255, 0.5*255, 0.5*255] = [127.5, 127.5, 127.5]
  • rknn_std = [0.5*255, 0.5*255, 0.5*255] = [127.5, 127.5, 127.5]

但请注意,RKNN的std_values参数在内部计算时是作为除数,即(x/255 - mean)/std。有些配置或代码习惯使用128作为std,这对应的是std=0.50196,与0.5有细微差别。务必保证这里的数值与训练时完全一致,一个数字的偏差都可能导致特征向量空间扭曲,比对失败。

2.2 构建转换脚本:量化与优化策略

下面是一个增强版的转换脚本convert_rknn.py,我加入了更多注释和实际开发中需要的环节。

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

# 模型路径
ONNX_MODEL = 'face_recog.onnx'
RKNN_MODEL = 'face_recog.rknn'

# 量化数据集生成函数
def create_dataset_txt(dataset_path, output_txt='./dataset.txt'):
    """
    生成量化所需的样本列表文件。
    假设dataset_path下存放着用于量化的图片(最好是来自验证集的人脸裁剪图)。
    """
    img_list = []
    for img_name in os.listdir(dataset_path):
        if img_name.endswith(('.jpg', '.png', '.jpeg')):
            img_path = os.path.join(dataset_path, img_name)
            # 这里可以添加与训练时一致的预处理,然后保存为二进制文件(.npy)供量化使用
            # 为简化,我们通常直接记录图片路径,RKNN工具会读取并预处理
            img_list.append(img_path + '\n')
    with open(output_txt, 'w') as f:
        f.writelines(img_list[:100]) # 量化不需要太多图片,通常100-200张足够

if __name__ == '__main__':
    # 1. 创建RKNN对象
    # verbose=True 可以打印更详细的日志,调试时很有用
    rknn = RKNN(verbose=True)

    # 2. 加载ONNX模型
    print('--> Loading ONNX model...')
    # 指定输出节点名,如果不确定,可以不指定,RKNN会尝试自动识别
    ret = rknn.load_onnx(model=ONNX_MODEL) #, outputs=['output_name'])
    if ret != 0:
        print('Load ONNX model failed!')
        exit(ret)
    print('Model loaded.')

    # 3. 配置模型
    print('--> Configuring model...')
    # target_platform: 必须指定为'rv1126'
    # mean_values/std_values: 根据你的训练预处理设置
    # quantized_dtype: 指定量化类型,'asymmetric_quantized-u8'是常用的INT8量化
    # quantized_algorithm: 量化算法,'normal'或'mmse'(最小均方误差),后者精度通常更好
    # quantized_method: 'channel'或'layer',channel粒度更细,精度更高
    ret = rknn.config(
        target_platform='rv1126',
        mean_values=[[127.5, 127.5, 127.5]], # 对应训练时 mean=[0.5,0.5,0.5]
        std_values=[[127.5, 127.5, 127.5]],   # 对应训练时 std=[0.5,0.5,0.5]
        quantized_dtype='asymmetric_quantized-u8',
        quantized_algorithm='mmse',
        quantized_method='channel',
        batch_size=1 # RV1126上batch_size=1是最常见的部署方式
    )
    if ret != 0:
        print('Config model failed!')
        exit(ret)
    print('Configuration done.')

    # 4. (可选)生成量化数据集
    # create_dataset_txt('./quant_images/', './dataset.txt')

    # 5. 构建模型(执行量化与编译)
    print('--> Building model...')
    start_time = time.time()
    ret = rknn.build(
        do_quantization=True,      # 开启量化
        dataset='./dataset.txt',   # 量化样本列表文件
        # pre_compile=True         # 预编译:生成的可执行文件在板端加载更快。但开启后模型将绑定特定NPU核心,泛化性降低。调试阶段建议先关闭。
    )
    build_time = time.time() - start_time
    print(f'Build time: {build_time:.2f} seconds')
    if ret != 0:
        print('Build model failed!')
        exit(ret)
    print('Model build successfully.')

    # 6. 导出RKNN模型文件
    print('--> Exporting RKNN model...')
    ret = rknn.export_rknn(RKNN_MODEL)
    if ret != 0:
        print('Export RKNN model failed!')
        exit(ret)
    print(f'RKNN model has been exported to: {RKNN_MODEL}')

    # 7. (重要)精度分析
    # 这一步会在PC上模拟NPU运行,并与浮点模型结果对比,给出量化后的精度损失评估
    print('--> Starting accuracy analysis...')
    ret = rknn.accuracy_analysis(inputs='./dataset.txt', target='rv1126')
    if ret != 0:
        print('Accuracy analysis failed!')
        # 不一定退出,但需要警惕精度损失过大
    print('Accuracy analysis done.')

    # 8. 释放资源
    rknn.release()

这个脚本比基础版本增加了几个关键部分:

  • 量化数据集准备:通过create_dataset_txt函数示意如何准备量化样本。这些样本应尽可能覆盖你应用场景的数据分布(不同光照、姿态的人脸),这对保持量化后精度至关重要。
  • 详细的配置参数:如quantized_algorithmquantized_method,它们直接影响最终的量化效果。mmse算法和channel方法通常能取得比默认设置更好的精度。
  • 预编译选项pre_compile是一个需要权衡的选项。开启后,模型在板端的加载速度会显著提升,但生成的模型文件与具体的NPU核心绑定,换一块板子可能就无法运行。建议在最终部署时才开启
  • 精度分析accuracy_analysis是一个极其有用的工具。它能生成一份报告,展示每一层量化前后的数值分布差异,帮助你定位是哪些层导致了大的精度损失,从而有针对性地调整量化策略或修改模型结构。

运行这个脚本,你就能得到可以在RV1126上加载的face_recog.rknn文件。

python convert_rknn.py

3. 模型测试与验证:确保转换无误

得到RKNN文件只是第一步,在部署到设备之前,必须在开发环境(通常是x86 Linux)中进行充分的模拟测试和精度验证。这个环节能帮你提前发现大部分问题,避免在板端进行低效的调试。

3.1 PC端模拟推理测试

编写一个test_on_pc.py脚本,使用RKNN的模拟器进行推理。这不需要真实的RV1126开发板。

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

def load_and_preprocess_image(image_path, input_size=(112, 112)):
    """模拟与训练时完全一致的预处理流程"""
    img = cv2.imread(image_path)
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) # 注意训练时可能使用RGB顺序
    img = cv2.resize(img, input_size)
    # 转换为float32并归一化到[0,1]
    img = img.astype(np.float32) / 255.0
    # 标准化 - 此处必须与转换配置完全一致!
    mean = np.array([0.5, 0.5, 0.5], dtype=np.float32)
    std = np.array([0.5, 0.5, 0.5], dtype=np.float32)
    img = (img - mean) / std
    # 调整维度顺序为 NCHW
    img = np.transpose(img, (2, 0, 1))
    img = np.expand_dims(img, axis=0) # 添加batch维度
    return img

def cosine_similarity(feat1, feat2):
    """计算两个特征向量的余弦相似度"""
    return np.dot(feat1, feat2) / (np.linalg.norm(feat1) * np.linalg.norm(feat2))

if __name__ == '__main__':
    rknn = RKNN()

    print('--> Load RKNN model')
    ret = rknn.load_rknn('./face_recog.rknn')
    if ret != 0:
        print('Load RKNN model failed')
        exit(ret)

    print('--> Init runtime (Simulator)')
    # 注意:这里target是'simulator',表示在PC上模拟运行
    ret = rknn.init_runtime(target='simulator')
    if ret != 0:
        print('Init runtime failed')
        exit(ret)

    # 测试单张图片
    test_img_path = './test_data/person_A_1.jpg'
    input_data = load_and_preprocess_image(test_img_path)

    print('--> Running inference')
    outputs = rknn.inference(inputs=[input_data])
    # outputs是一个列表,获取第一个输出(特征向量)
    feature = outputs[0]
    print(f'Feature shape: {feature.shape}')
    print(f'Feature (first 10 dims): {feature[0][:10]}')

    # 测试人脸比对
    img1_path = './test_data/person_A_1.jpg'
    img2_path = './test_data/person_A_2.jpg' # 同一个人的另一张照片
    img3_path = './test_data/person_B_1.jpg' # 另一个人的照片

    feat1 = rknn.inference(inputs=[load_and_preprocess_image(img1_path)])[0]
    feat2 = rknn.inference(inputs=[load_and_preprocess_image(img2_path)])[0]
    feat3 = rknn.inference(inputs=[load_and_preprocess_image(img3_path)])[0]

    # 对特征进行L2归一化(如果模型输出未归一化)
    feat1_norm = feat1 / np.linalg.norm(feat1)
    feat2_norm = feat2 / np.linalg.norm(feat2)
    feat3_norm = feat3 / np.linalg.norm(feat3)

    sim_same = cosine_similarity(feat1_norm[0], feat2_norm[0])
    sim_diff = cosine_similarity(feat1_norm[0], feat3_norm[0])

    print(f'\n--- Similarity Results ---')
    print(f'Same person (A1 vs A2): {sim_same:.4f}')
    print(f'Different person (A1 vs B1): {sim_diff:.4f}')

    # 可以根据经验设定一个阈值,例如0.5或0.6
    threshold = 0.5
    print(f'\n--- Verification (Threshold={threshold}) ---')
    print(f'A1 vs A2: {"PASS" if sim_same > threshold else "FAIL"}')
    print(f'A1 vs B1: {"PASS" if sim_diff < threshold else "FAIL"}')

    rknn.release()

这个脚本的关键在于:

  1. 完全复现预处理load_and_preprocess_image函数必须与训练和转换时的预处理逻辑逐字节对齐。
  2. 使用模拟器target='simulator'让你在没有硬件的情况下验证模型功能。
  3. 进行定量测试:通过计算同一人不同照片、以及不同人照片之间的特征相似度,可以直观地评估模型转换后的区分能力是否保持。如果sim_same值骤降或sim_diff值飙升,说明预处理或量化出了问题。

3.2 与原始浮点模型对比

最可靠的验证方法,是将RKNN模型的输出与原始PyTorch/TensorFlow浮点模型在同一批测试数据上的输出进行对比。计算它们输出的特征向量之间的余弦相似度或欧氏距离。理想情况下,这个差异应该非常小(例如余弦相似度>0.999)。如果差异明显,就需要回溯检查预处理、量化配置甚至模型导出环节。

4. RV1126板端部署与性能调优

当PC端测试通过后,就可以将模型部署到真实的RV1126开发板上了。这一步的目标是确保模型能在资源受限的嵌入式环境中稳定、高效地运行。

4.1 交叉编译与环境部署

首先,需要在你的开发主机上,使用RV1126的SDK中的交叉编译工具链,编译出板端可执行的C++或Python推理程序。瑞芯微官方通常提供示例代码。这里以Python API为例,因为它更便于快速原型验证。

将以下文件传输到RV1126开发板:

  • face_recog.rknn 模型文件
  • 测试图片
  • 一个精简版的Python推理脚本(基于RKNN的Python API,板端需要安装对应的RKNN Runtime库)。

板端推理脚本inference_on_device.py与PC测试脚本类似,但init_runtime的target需要改为'rv1126',并且要注意内存管理。

# inference_on_device.py (在RV1126上运行)
import numpy as np
import cv2
from rknnlite.api import RKNNLite # 注意:板端可能使用RKNNLite接口,更轻量

def run_inference():
    rknn_lite = RKNNLite()

    # 加载模型
    ret = rknn_lite.load_rknn('./face_recog.rknn')
    if ret != 0:
        print('Load RKNN model failed')
        return

    # 初始化运行时环境,指定RKNN核心编号(如果有多核)
    ret = rknn_lite.init_runtime(core_mask=RKNNLite.NPU_CORE_0)
    if ret != 0:
        print('Init runtime failed')
        return

    # ... 图像加载和预处理(与PC端完全一致)...
    # ... 执行推理 rknn_lite.inference() ...
    # ... 后处理 ...

    rknn_lite.release()

if __name__ == '__main__':
    run_inference()

4.2 性能分析与瓶颈定位

在板端运行模型后,你需要关注几个关键性能指标:

  • 推理耗时:单次前向传播的时间。使用time.time()inference()调用前后计时。
  • 内存占用:模型加载和运行时占用的内存。可以通过板端的free命令或RKNN提供的接口查看。
  • 功耗:对于电池供电的设备尤为重要。

如果性能不达标,可以考虑以下优化方向:

  • 调整模型输入尺寸:这是最有效的优化手段之一。将输入从112x112降到96x96甚至80x80,能显著减少计算量和内存带宽消耗,但可能会牺牲一些精度。需要做权衡测试。
  • 使用预编译模型:在最终生产固件中,使用pre_compile=True生成的RKNN模型,可以避免板端的模型解析和编译时间,加快首次加载速度。
  • 启用NPU多核:如果模型较大,可以尝试将计算图分割到多个NPU核心上并行执行。在init_runtime时设置core_mask=RKNNLite.NPU_CORE_0_1
  • 优化前后处理:图像预处理(缩放、归一化)和后处理(特征归一化、比对)如果在CPU上进行,也可能成为瓶颈。考虑使用OpenCV的优化,或者将部分操作(如减均值)通过RKNN的mean_values参数卸载到NPU内部完成。

4.3 常见问题与调试技巧

在板端部署时,你可能会遇到一些典型问题:

  • 模型加载失败:检查模型文件路径是否正确,板端RKNN Runtime库版本是否与转换工具版本匹配。
  • 推理结果异常:首先确认板端预处理代码与PC测试代码绝对一致。检查图像通道顺序(BGR vs RGB)、数据精度(float32 vs uint8)。可以在板端和PC端对同一张图片运行预处理,并对比处理后的原始数组数据,确保完全相同。
  • 内存不足:RV1126的可用内存有限。如果模型太大或同时运行多个任务,可能导致分配失败。尝试减小模型尺寸、降低输入分辨率或关闭其他进程。
  • 精度下降:如果PC模拟结果好,但板端精度差,可能是量化过程中某些层对硬件不友好,导致数值误差放大。可以尝试:
    • rknn.config中调整quantized_algorithm'mmse'
    • 使用更多样、更接近真实场景的图片作为量化数据集。
    • 对精度损失特别大的层,尝试在rknn.build中通过custom_layer参数将其设置为不量化(保持浮点),但这可能会影响性能。

整个从ONNX到RKNN的转换和部署流程,是一个需要耐心调试和反复验证的工程。它要求开发者不仅理解算法,还要熟悉硬件特性和工具链的细节。我最深的体会是,建立一套自动化的验证流水线至关重要——从模型导出、转换、PC测试到板端验证,每一步都应有脚本和标准数据,确保任何改动都能被快速、准确地评估。在RV1126上成功跑通一个人脸识别模型,那种成就感远大于在服务器上训练出一个新SOTA,因为这意味着你的算法真正走出了实验室,开始在现实世界中发挥作用。

Logo

智能硬件社区聚焦AI智能硬件技术生态,汇聚嵌入式AI、物联网硬件开发者,打造交流分享平台,同步全国赛事资讯、开展 OPC 核心人才招募,助力技术落地与开发者成长。

更多推荐