4、RKNN输入格式要求:模型转换时如何指定输入图像的格式和Layout

在RKNN模型部署流程中,模型转换(rknn-toolkit2或rknn-toolkit的convert阶段)是决定推理性能与精度的关键环节。其中,输入图像的格式(Format)与数据排布(Layout)直接决定了NPU对张量的解析方式。如果转换时指定的输入格式与后续预处理代码或实际采集的图像数据不一致,轻则导致推理结果错乱,重则引发内存越界或驱动报错。本章将系统讲解RKNN模型转换时输入图像格式与Layout的指定方法、底层原理、判断流程及实战案例。

RKNN模型转换输入格式与Layout决策链路 输入源 摄像头/图像/视频 BGR/RGB/NV12 预处理 resize/crop/归一化 通道顺序转换 rknn.config() mean_values/std_values quant_img_RGB2BGR RKNN 模型 Layout 决策 NHWC 高度×宽度×通道 NCHW 通道×高度×宽度 RKNN默认NHWC,ONNX转RKNN时 自动插入Transpose处理 建议:统一使用NHWC + RGB + 0~255 格式决策 RGB / BGR 三通道彩色图像 NV12 / NV21 YUV420半平面 RKNN支持RGB、BGR、GRAY、NV12、NV21 量化数据集需与推理输入格式一致 注意:量化图片与推理图片格式必须统一

4.1 核心原理:RKNN输入张量的数据描述

RKNN NPU内部采用固定维度的张量计算引擎,其输入张量描述由dtype(数据类型)、format(颜色格式)和layout(数据排布)三个维度共同定义。在模型转换阶段,rknn-toolkit2通过读取原始模型(如ONNX、Caffe、TensorFlow)的输入节点信息,结合用户在rknn.config()和rknn.load_input()中指定的参数,生成一个内部统一的张量描述。

其中,Layout指的是数据在内存中的排列顺序。RKNN默认使用NHWC(Batch × Height × Width × Channel)排布,这与TensorFlow的原生格式一致。而ONNX模型通常使用NCHW(Batch × Channel × Height × Width)。当用户加载ONNX模型时,rknn-toolkit2会自动在输入节点前插入一个Transpose算子,将NCHW转换为NHWC。但这一自动转换并非总是最优——如果原始模型已经包含了对输入数据的特殊处理(例如某些自定义算子直接操作NCHW数据),则需要在转换时显式指定layout,避免重复转换导致性能损失。

Format则描述的是图像通道顺序与色彩空间。RKNN支持RGB、BGR、GRAY(单通道灰度)、NV12、NV21等格式。其中NV12/NV21是YUV420格式的两种半平面变体,常用于摄像头直接输出。在转换时,如果指定了quant_img_RGB2BGR=True,则工具会将输入的RGB量化图片转换为BGR顺序后再进行量化校准,从而保证量化统计的均值/方差与推理时输入数据分布一致。

关键理解:格式(Format)解决“通道是什么颜色”的问题,Layout解决“通道在内存中怎么排列”的问题。两者独立但必须同时正确,否则NPU读取到的数据就是“错位”的。

4.2 判断方法与操作流程:转换时如何指定

在rknn-toolkit2的Python API中,指定输入格式和Layout主要涉及以下三个环节:

  1. rknn.config():设置全局量化与预处理参数,包括mean_values、std_values、quant_img_RGB2BGR、target_platform等。
  2. rknn.load_input():加载用于量化校准的图片数据集,此时需要确保图片的通道顺序与config中设置一致。
  3. rknn.build():构建RKNN模型,此阶段会根据原始模型的输入节点自动推断Layout,但用户可以通过rknn.config(optimization_level=3)等参数影响布局优化策略。

对于Layout的显式控制,rknn-toolkit2提供了rknn.config(reorder_channel='0 1 2')参数,用于手动指定输入通道的排列顺序。例如,如果原始模型输入是BGR顺序,而量化图片是RGB,则可以通过reorder_channel='2 1 0'进行通道重排。但更常见的做法是使用quant_img_RGB2BGR参数,让工具自动处理。

以下是一个典型的转换配置代码片段:

from rknn.api import RKNN

rknn = RKNN()
# 配置量化与预处理参数
rknn.config(
    mean_values=[[123.675, 116.28, 103.53]],  # 与ImageNet一致的均值
    std_values=[[58.395, 57.12, 57.375]],
    quant_img_RGB2BGR=True,  # 量化图片为RGB,但模型输入为BGR
    target_platform='rk3588'
)
# 加载ONNX模型
rknn.load_onnx(model='yolov5s.onnx')
# 加载量化图片(图片为RGB格式)
rknn.load_input(input_list='dataset.txt')
# 构建RKNN模型
rknn.build(do_quantization=True, dataset='dataset.txt')
rknn.export_rknn('yolov5s.rknn')

在上述代码中,quant_img_RGB2BGR=True告诉工具:量化图片是RGB顺序,但模型实际期望BGR输入,因此工具会在量化统计时自动将图片转换为BGR。这样,在推理阶段,用户只需将预处理后的BGR数据直接送入NPU即可,无需额外转换。

4.3 案例拆解:YOLOv5s从ONNX到RKNN的格式适配

以RK3588平台部署YOLOv5s为例,原始ONNX模型的输入节点为images,形状为[1, 3, 640, 640](NCHW),数据类型为float32。在转换过程中,我们遇到以下问题:

  • 初始现象:直接使用默认配置转换后,在开发板上推理输出大量错误的检测框,且置信度普遍低于0.1。
  • 分析过程:通过打印RKNN模型的输入张量信息,发现工具自动插入了Transpose算子将NCHW转为NHWC,但量化统计时使用的图片是RGB顺序,而YOLOv5s训练时输入是RGB,但ONNX导出时部分版本会转换为BGR。这导致量化均值和实际推理数据分布不一致。
  • 关键数据:原始模型mean=[0,0,0],std=[1,1,1](即未归一化),但量化图片像素范围0~255。若不设置mean/std,量化误差会显著增大。

解决步骤如下:

  1. 确认ONNX模型输入Layout为NCHW,但RKNN内部使用NHWC,因此无需手动干预,工具自动处理。
  2. 在rknn.config()中设置mean_values=[[0,0,0]]、std_values=[[255,255,255]],将像素范围从0~255归一化到0~1。
  3. 设置quant_img_RGB2BGR=False,因为YOLOv5s的ONNX输入为RGB顺序,量化图片也使用RGB,保持一致。
  4. 重新构建模型后,在开发板上使用相同预处理(resize到640×640,保持RGB顺序,除以255)进行推理,检测结果恢复正常。

此案例表明,格式不一致是导致量化模型精度暴跌的首要原因。在转换前,务必通过Netron等工具查看原始模型的输入节点信息,确认其Layout和通道顺序。

风险提示:如果原始模型输入是NCHW,而你在推理代码中手动将数据reshape为NHWC并直接送入NPU,会导致数据完全错乱。RKNN的NPU驱动期望输入张量严格遵循模型转换时确定的Layout,任何手动reshape操作都是多余的,且可能引发内存访问越界。

4.4 常见误区与总结复盘

误区 后果 正确做法
忽略ONNX的NCHW与RKNN的NHWC差异 推理结果完全错误,或性能下降 依赖工具自动Transpose,不要手动干预
量化图片与推理图片通道顺序不一致 量化统计偏差,精度下降 统一使用RGB或BGR,并用quant_img_RGB2BGR对齐
mean/std设置与训练时不一致 输入分布偏移,检测率下降 从训练脚本中复制mean/std值
使用NV12输入但量化图片用RGB 量化统计失效,模型输出噪声 量化图片需与推理输入格式完全一致

总结复盘:在RKNN模型转换中,输入格式与Layout的指定并非“填个参数”那么简单,而是需要从原始模型定义出发,结合量化数据集和推理端预处理代码,形成一条完整的数据链路一致性。建议在项目初期就制定一份输入数据规范文档,明确以下内容:

  • 原始模型的输入Layout(NCHW还是NHWC)与通道顺序(RGB/BGR)
  • 量化图片的存储格式与读取方式
  • 推理端预处理步骤(resize、归一化、通道转换)
  • RKNN config中对应的mean/std/quant_img_RGB2BGR设置