在海思SD3403板子上跑通YOLOv5后,尝试将YOLOv8部署到3403上,目前全流程已跑通,算法模型可以在板端的NNN和SVP_NNN核上稳定运行。
首先话不多说进入YOLOv8模型训练部分,不要使用pip install ultralytics安装库,而是下载源码,对源码中训练部分进行修改。这里就不挂链接了,YOLOv8原版模型是无法直接部署到板子上的(至少在我的pipeline里无法直接部署),原因是训练好的pytorch模型需要转换成onnx模型再转换为caffemodel,caffe目前不支持Upsample层和sub层。
训练YOLOv8模型需要对模型定义文件(ultralytics/cfg/models/v8/yolov8.yaml)中head部分进行改造。
# YOLOv8.0n head
head:
- [-1, 1, nn.Upsample, [None, 2, "nearest"]]
- [[-1, 6], 1, Concat, [1]] # cat backbone P4
- [-1, 3, C2f, [512]] # 12
- [-1, 1, nn.Upsample, [None, 2, "nearest"]]
- [[-1, 4], 1, Concat, [1]] # cat backbone P3
- [-1, 3, C2f, [256]] # 15 (P3/8-small)
- [-1, 1, Conv, [256, 3, 2]]
- [[-1, 12], 1, Concat, [1]] # cat head P4
- [-1, 3, C2f, [512]] # 18 (P4/16-medium)
- [-1, 1, Conv, [512, 3, 2]]
- [[-1, 9], 1, Concat, [1]] # cat head P5
- [-1, 3, C2f, [1024]] # 21 (P5/32-large)
- [[15, 18, 21], 1, Detect, [nc]] #