YOLOv5-7.0 修改源码到导出到RKNN支持的onnx格式导出指南

YOLOv5-7.0 RKNN导出指南

概述

本文档说明如何将YOLOv5模型导出为适合RKNN（Rockchip Neural Network）推理的ONNX格式。

修改内容

1. models/yolo.py 修改

为了适配RKNN推理，修改了Detect类的forward方法：

原始代码（已注释）：

def forward(self, x):
    z = []  # inference output
    for i in range(self.nl):
        x[i] = self.m[i](x[i])  # conv
        bs, _, ny, nx = x[i].shape  # x(bs,255,20,20) to x(bs,3,20,20,85)
        x[i] = x[i].view(bs, self.na, self.no, ny, nx).permute(0, 1, 3, 4, 2).contiguous()
        # ... 复杂的后处理逻辑
    return x if self.training else (torch.cat(z, 1),) if self.export else (torch.cat(z, 1), x)

RKNN适配代码（当前使用）：

def forward(self, x): #修改，为了导出rknn适配模式
    z = []  # inference output
    for i in range(self.nl):
        x[i] = self.m[i](x[i])  # conv
    return x

2. export.py 修改

修改了输出形状处理逻辑，以支持RKNN模式下的多输出格式。

修改位置: export.py 第552行

完整的修改前后对比:

修改前 (原始代码):

    for _ in range(2):
        y = model(im)  # dry runs
    if half and not coreml:
        im, model = im.half(), model.half()  # to FP16
    shape = tuple((y[0] if isinstance(y, tuple) else y).shape)  # model output shape
    metadata = {'stride': int(max(model.stride)), 'names': model.names}  # model metadata
    LOGGER.info(f"\n{colorstr('PyTorch:')} starting from {file} with output shape {shape} ({file_size(file):.1f} MB)")

修改后 (当前代码):

    for _ in range(2):
        y = model(im)  # dry runs
    if half and not coreml:
        im, model = im.half(), model.half()  # to FP16
    # Handle different output formats for RKNN export
    if isinstance(y, list):
        # For RKNN export mode, y is a list of tensors
        shape = [tuple(tensor.shape) for tensor in y]  # list of shapes for each output
        LOGGER.info(f"Model outputs {len(y)} tensors with shapes: {shape}")
    else:
        shape = tuple((y[0] if isinstance(y, tuple) else y).shape)  # model output shape
    metadata = {'stride': int(max(model.stride)), 'names': model.names}  # model metadata
    LOGGER.info(f"\n{colorstr('PyTorch:')} starting from {file} with output shape {shape} ({file_size(file):.1f} MB)")

具体修改内容:

1. 删除了: shape = tuple((y[0] if isinstance(y, tuple) else y).shape) # model output shape
2. 添加了:
   • 注释: # Handle different output formats for RKNN export
   • 条件判断: if isinstance(y, list):
   • list处理: shape = [tuple(tensor.shape) for tensor in y]
   • 调试信息: LOGGER.info(f"Model outputs {len(y)} tensors with shapes: {shape}")
   • else分支: 保留原始的tuple/tensor处理逻辑

修改原因:

• 问题: 原始代码假设模型输出是tuple或单个tensor，无法处理list类型
• 错误: RKNN模式下会报错 AttributeError: 'list' object has no attribute 'shape'
• 解决: 新代码能够正确处理list、tuple和tensor三种输出格式
• 兼容性: 保持了对原始输出格式的兼容性

导出命令

基本导出命令

python export.py --weights yolov5s.pt --img 640 --batch 1 --include onnx

参数说明

• --weights: 模型权重文件路径
• --img: 输入图像尺寸
• --batch: 批处理大小（RKNN通常使用1）
• --include onnx: 导出ONNX格式

输出格式

RKNN模式输出

模型输出3个不同尺度的特征图：

• 输出1: (1, 255, 80, 80) - 8倍下采样
• 输出2: (1, 255, 40, 40) - 16倍下采样
• 输出3: (1, 255, 20, 20) - 32倍下采样

每个输出的通道数255 = (类别数80 + 5) × 锚点数3

与标准YOLOv5的区别

1. 标准模式: 输出单个张量 (1, 25200, 85)，包含所有检测框
2. RKNN模式: 输出3个原始特征图，后处理在RKNN推理时进行

后续步骤

转换为RKNN模型

使用导出的ONNX文件转换为RKNN格式：

convert.py文件在rknn_model_zoo-main/examples/yolov5/python下，rknn_model_zoo-main是官网github下载的；

#：浮点，不量化【选这个】 
python convert.py ../model/yolov5s.onnx rk3588 fp


# 转换为 RK3588 量化模型
python convert.py ../model/yolov5s.onnx rk3588 i8 ../model/yolov5s_rk3588.rknn

# 转换为 RK3588 浮点模型
python convert.py ../model/yolov5s.onnx rk3588 fp ../model/yolov5s_rk3588_fp.rknn

# 转换为 RK3566 模型
python convert.py ../model/yolov5s.onnx rk3566

支持的平台

• RK3562, RK3566, RK3568, RK3576, RK3588
• RV1103, RV1106, RV1126B, RV1109, RV1126
• RK1808, RK3399PRO

注意事项

1. 版本兼容性: 确保PyTorch和torchvision版本兼容
2. 后处理: RKNN模式下，NMS等后处理需要在推理端实现
3. 精度: 量化可能影响模型精度，建议测试验证
4. 内存: RKNN推理时注意内存使用优化

恢复标准模式

如需恢复标准YOLOv5模式，取消注释models/yolo.py中的原始forward方法，并注释掉RKNN适配代码。

文件清单

• yolov5s.onnx: 导出的ONNX模型文件
• models/yolo.py: 修改的模型定义文件
• export.py: 修改的导出脚本
• RKNN_Export_Guide.md: 本文档