4. 模型转换前置准备

在RKNN工具链的部署流程中，ONNX（Open Neural Network Exchange）模型是连接PyTorch/TensorFlow训练框架与RKNN编译器之间的关键中间表示。RKNN-Toolkit2虽然支持直接导入PyTorch模型，但在实际工程中，我们强烈建议先导出为ONNX格式，再进行RKNN转换。这是因为ONNX具备标准化的计算图描述、完善的算子集定义以及丰富的模型验证工具链，能够显著降低因框架版本差异、算子映射缺失导致的转换失败风险。本章将系统讲解导出ONNX时的输入尺寸策略（动态/静态）、NMS（Non-Maximum Suppression，非极大值抑制）层的处理决策，以及如何通过可视化与数值比对验证ONNX模型的正确性。

上图展示了从PyTorch模型到RKNN转换前的完整准备链路。三个关键阶段环环相扣：导出配置决定了ONNX模型的输入灵活性与编译效率；NMS处理策略直接影响后处理代码的复杂度与NPU利用率；正确性验证则是确保后续RKNN转换结果可信的基石。下面我们将逐一深入剖析每个环节的技术细节与工程决策。

4.1 导出ONNX模型：输入尺寸的动态与静态抉择

4.1.1 核心原理：静态尺寸与动态尺寸的本质差异

ONNX模型的输入尺寸定义方式直接影响RKNN编译器对模型的计算图优化程度。静态尺寸（Static Shape）在导出时将所有输入维度固定为具体数值，例如[1, 3, 640, 640]。RKNN编译器可以基于固定形状进行内存规划、算子融合和通道重排等深度优化，生成最精简的NPU指令序列。而动态尺寸（Dynamic Shape）允许在推理时改变batch大小或图像宽高，通过dynamic_axes参数指定可变维度。虽然动态尺寸提升了模型灵活性，但RKNN工具链对动态形状的支持存在诸多限制：动态维度会导致NPU无法预分配静态内存，可能触发多次内存重映射，显著增加推理延迟；同时，部分RKNN算子（如Resize、Concat）在动态形状下无法进行常量折叠优化。

对比维度	静态尺寸（推荐）	动态尺寸
RKNN编译优化	深度优化，算子融合充分	优化受限，存在额外开销
推理延迟	最低，适合实时场景	较高，动态内存分配耗时
部署灵活性	固定输入分辨率，需预处理resize	可适应多分辨率输入
RKNN支持度	完全支持	部分支持，需查阅版本说明
适用场景	固定摄像头分辨率、边缘盒子	多尺寸输入、动态batch服务

4.1.2 操作流程：导出静态尺寸ONNX模型

以YOLOv8为例，推荐使用静态尺寸导出。以下是标准导出代码片段：

import torch
from ultralytics import YOLO

# 加载训练好的模型
model = YOLO("runs/detect/train/weights/best.pt")
model.model.eval()

# 构造静态输入张量 [batch, channel, height, width]
dummy_input = torch.randn(1, 3, 640, 640)

# 导出ONNX，opset_version=12 兼容RKNN
torch.onnx.export(
    model.model,
    dummy_input,
    "yolov8n.onnx",
    opset_version=12,
    input_names=["images"],
    output_names=["output0"],
    dynamic_axes=None  # 关键：不设置动态轴
)
print("静态ONNX导出完成")

若确需动态尺寸，例如需要支持不同分辨率的输入，可指定动态轴：

dynamic_axes = {
    "images": {0: "batch", 2: "height", 3: "width"},
    "output0": {0: "batch", 2: "anchors"}
}
torch.onnx.export(..., dynamic_axes=dynamic_axes)

4.1.3 案例拆解：从动态尺寸回退到静态尺寸的调优

背景：某安防项目需要在RK3588平台上部署YOLOv8s模型，原始PyTorch模型支持动态输入，用于处理不同分辨率的摄像头画面（1280×720、1920×1080）。初始现象：直接使用动态尺寸导出ONNX并转换RKNN后，模型推理帧率仅为12 FPS，且内存占用波动明显。分析过程：通过RKNN性能分析工具发现，动态Resize算子导致NPU流水线频繁停顿，且每次推理时内存重映射耗时约8ms。实际操作：将模型固定为1280×1280静态输入（统一resize），重新导出ONNX并转换。推理帧率提升至28 FPS，内存占用稳定。结果对比：虽然预处理增加了resize计算，但整体端到端延迟从83ms降至36ms。复盘经验：在边缘部署中，固定输入分辨率带来的性能收益远大于灵活性收益，应优先考虑静态尺寸。

4.2 处理NMS层：导出时去除还是保留

4.2.1 核心原理：NMS在ONNX中的位置与RKNN的算子支持

NMS是目标检测后处理的核心步骤，用于消除重叠的检测框。在PyTorch模型中，NMS通常以自定义函数或TorchScript脚本形式存在，不属于标准ONNX算子集。当使用torch.onnx.export导出时，NMS可能被转换为多个基础算子组合（如Sort、Gather、NonZero等），导致计算图膨胀且难以优化。更关键的是，RKNN NPU对NMS这类非规则控制流算子的支持非常有限——RKNN-Toolkit2虽然提供了NonMaxSuppression自定义算子，但其实现仅支持特定输入格式，且性能不如CPU上的优化实现。

4.2.2 判断方法与推荐策略

根据YOLO系列模型的结构差异，我们推荐以下处理策略：

• YOLOv5/v8（Anchor-Free）：导出时直接去除NMS，输出裸张量（形状为[1, 84, 8400]或[1, 4+80, 8400]）。后处理（解码+置信度过滤+NMS）在RKNN的CPU端或独立NPU推理后处理单元中完成。
• YOLOv3/v5（Anchor-Based）：同样去除NMS，输出三个尺度的原始预测张量，在外部实现解码与NMS。
• YOLOX：去除NMS，输出[1, 8400, 85]的原始预测。

4.2.3 实战场景：YOLOv8后处理代码框架

去除NMS后，ONNX输出为[1, 84, 8400]，其中84 = 4（框坐标）+ 80（COCO类别概率）。后处理流程如下：

import numpy as np
import cv2

def postprocess(output, conf_thres=0.25, iou_thres=0.45):
    # output shape: [1, 84, 8400] -> 转置为 [8400, 84]
    preds = output[0].transpose(1, 0)  # [8400, 84]
    boxes = preds[:, :4]  # cx, cy, w, h
    class_scores = preds[:, 4:]  # 80个类别概率
    # 解码框坐标（YOLOv8格式）
    boxes_xyxy = np.zeros_like(boxes)
    boxes_xyxy[:, 0] = boxes[:, 0] - boxes[:, 2] / 2
    boxes_xyxy[:, 1] = boxes[:, 1] - boxes[:, 3] / 2
    boxes_xyxy[:, 2] = boxes[:, 0] + boxes[:, 2] / 2
    boxes_xyxy[:, 3] = boxes[:, 1] + boxes[:, 3] / 2
    # 获取最高类别分数
    class_ids = np.argmax(class_scores, axis=1)
    confidences = class_scores[np.arange(len(class_ids)), class_ids]
    # 置信度过滤
    mask = confidences > conf_thres
    boxes, scores, class_ids = boxes_xyxy[mask], confidences[mask], class_ids[mask]
    # NMS
    indices = cv2.dnn.NMSBoxes(boxes.tolist(), scores.tolist(), conf_thres, iou_thres)
    return boxes[indices], scores[indices], class_ids[indices]

4.3 验证ONNX正确性：结构检查与数值比对

4.3.1 核心原理：双重验证机制

ONNX模型验证包含两个层面：结构验证确保计算图符合ONNX规范，算子连接正确；数值验证确保ONNX模型的推理输出与原始PyTorch模型一致。结构验证通过onnx.checker完成，数值验证则借助ONNX Runtime（ORT）执行推理并与PyTorch输出对比。由于RKNN转换过程会进一步优化计算图，如果ONNX阶段就存在数值偏差，RKNN阶段将无法纠正，因此ONNX验证是质量保障的第一道防线。

4.3.2 操作流程：完整的验证步骤

第一步：结构验证

import onnx
model = onnx.load("yolov8n.onnx")
onnx.checker.check_model(model)
print("ONNX结构检查通过")

第二步：Netron可视化检查使用Netron工具打开ONNX文件，重点检查：输入节点名称与维度、输出节点名称与维度、是否有异常算子（如自定义NMS节点）、数据流是否连贯。

第三步：数值比对

import onnxruntime as ort
import torch
import numpy as np

# 准备测试输入
test_input = np.random.randn(1, 3, 640, 640).astype(np.float32)

# PyTorch推理
with torch.no_grad():
    torch_output = model.model(torch.from_numpy(test_input)).cpu().numpy()

# ONNX Runtime推理
ort_session = ort.InferenceSession("yolov8n.onnx")
ort_output = ort_session.run(None, {"images": test_input})[0]

# 数值比对
cos_sim = np.dot(torch_output.flatten(), ort_output.flatten()) / (
    np.linalg.norm(torch_output) * np.linalg.norm(ort_output))
max_diff = np.max(np.abs(torch_output - ort_output))
print(f"余弦相似度: {cos_sim:.6f}")
print(f"最大误差: {max_diff:.6e}")

验证指标	通过标准	失败排查方向
余弦相似度	> 0.99	检查模型是否处于eval模式、BN层是否冻结
最大误差	< 1e-4	检查opset版本、算子精度（FP16 vs FP32）
输出形状	完全一致	检查dynamic_axes设置、输出张量顺序
结构检查	无错误	检查是否有未支持的算子、图是否断开

4.3.3 案例拆解：数值偏差的定位与修复

背景：某自动驾驶团队在导出YOLOv5s ONNX模型后，发现ORT输出与PyTorch输出存在显著差异，余弦相似度仅为0.85。初始现象：模型在PyTorch中检测正常，但ONNX模型在RKNN上完全无法检测到目标。分析过程：首先进行数值比对，发现最大误差达到0.5，远超阈值。进一步检查发现，模型导出时未调用model.eval()，导致BatchNorm层仍处于训练模式，其运行均值/方差未固定。实际操作：在导出前显式调用model.eval()，并设置torch.no_grad()上下文。重新导出后，余弦相似度提升至0.9999，最大误差降至1e-5。结果对比：RKNN转换后模型检测精度与PyTorch原始模型完全一致。复盘经验：导出ONNX前必须确保模型处于推理模式，且所有随机性操作（如Dropout）已禁用。

4.4 常见误区与总结复盘

4.4.1 高频错误清单

• 误区一：盲目使用动态尺寸。许多开发者为了“通用性”设置动态宽高，结果在RKNN转换时遭遇“Unsupported dynamic shape”错误，或推理性能骤降。
• 误区二：保留NMS节点。部分教程建议保留NMS以简化后处理，但RKNN对NMS的支持不成熟，且NPU执行NMS效率低下。
• 误区三：忽略opset版本。RKNN-Toolkit2对ONNX opset版本有明确要求（通常支持opset 10-12），使用opset 13+可能导致算子不兼容。
• 误区四：验证不充分。仅检查ONNX能否加载，未进行数值比对，导致RKNN转换后精度异常时难以定位是ONNX问题还是RKNN量化问题。

4.4.2 最佳实践总结

综合本章内容，模型转换前置准备的最佳实践可归纳为：固定静态输入尺寸（如640×640）、导出时去除NMS层、使用opset 12、通过Netron和数值比对双重验证。这一流程能最大化RKNN编译器的优化空间，同时保证模型精度无损传递。在实际项目中，建议将ONNX导出与验证脚本固化为CI流程的一部分，每次模型更新后自动执行验证，确保进入RKNN转换阶段的模型质量始终可控。