YOLOv5 移植 RK3588 踩坑记录

YOLOv5 移植 RK3588 踩坑记录

（pt → onnx → rknn 全流程）

环境：YOLOv5-5.x / rknn-toolkit2 / Python≥3.8 板卡：RK3588
目标：复用官方 yolov5_rknn 示例，仅替换 *.rknn + *.yaml 即可跑通

---

1. pt → onnx 转换

1.1 原生 export.py 直接转会报错

python export.py --weights yolov5s.pt --include onnx

运行 rknn 示例 test.py 时：

IndexError: list index out of range

原因：Detect 层输出 tuple，而 rknn 示例只拿 list 第 1 个特征图。

---

1.2 修改 1：去掉 Detect 后处理，只输出 3 个 sigmoid 特征图

文件 models/yolo.py
类 Detect.forward()
原始代码（已注释）：

def forward(self, x):
    z = []
    for i in range(self.nl):
        x[i] = self.m[i](x[i])  # conv
        bs, _, ny, nx = x[i].shape
        x[i] = x[i].view(bs, self.na, self.no, ny, nx).permute(0, 1, 3, 4, 2).contiguous()
        if not self.training:
            ...
            z.append(y.view(bs, self.na * nx * ny, self.no))
    return x if self.training else (torch.cat(z, 1), ) if self.export else (torch.cat(z, 1), x)

修改为：

def forward(self, x):
    z = []
    for i in range(self.nl):
        z.append(torch.sigmoid(self.m[i](x[i])))  # 仅卷积+sigmoid
    return z

说明：

1. 不再做 view/permute/grid/anchor 等后处理；
2. 返回 List[Tensor]，rknn 侧直接拿 3 个特征图。

---

1.3 修改 2：export.py 支持 list 类型输出

文件 export.py
位置：run() 函数里获取 output shape 处
原始：

shape = tuple((y[0] if isinstance(y, tuple) else y).shape)

改为：

shape = tuple((y[0] if (isinstance(y, tuple) or isinstance(y, list)) else y).shape)

保证 y 为 list 时也能正常取第 0 个元素。

---

1.4 重新转 onnx

python export.py --weights yolov5s.pt --include onnx

用 Netron 打开，确认尾部只剩 3 个 Sigmoid 节点，无 Detect 算子。

---

2. anchors 文件生成（踩坑 2）

背景：RKNN 示例后处理需要 anchors，必须与训练时保持一致。
工具：官方 yolov5_rknn 已提供 export.py，无需手动改 anchor。

git clone https://github.com/airockchip/yolov5_rknn
cd yolov5_rknn
python export.py --rknpu --weights car_tk_hanma.pt

执行完毕会自动生成：

RK_anchors.txt

内容示例（9 行）：

32.0
32.0
...

若用自己改过的 export.py 忘记生成，会导致 rknn 推理框全部漂移或 0 检出。

---

3. onnx → rknn

1. 把上一步得到的 *.onnx + RK_anchors.txt 复制到
   rknn-toolkit2/examples/yolov5/python 目录。

2. 修改 yaml 文件（类别名/anchor 路径）：

   anchors_path: RK_anchors.txt
   names: ['class1', 'class2', ...]
   

3. 执行：

   python test.py
   

   会自动完成 onnx → rknn 转换 + 板端推理。

---

4. 一键总结

步骤	关键修改	文件
① 去掉 Detect 后处理	forward 返回 List[sigmoid(feat)]	models/yolo.py
② export 支持 list	isinstance(y, tuple) or isinstance(y, list)	export.py
③ 生成 anchors	加 --rknpu 开关	官方 export.py
④ 转 rknn	直接跑 test.py	rknn-toolkit2 示例

---