YOLOV5从模型训练到部署在NPU上实现流程

原创

已于 2024-03-27 11:06:55 修改 · 2.7k 阅读

25 ·

CC 4.0 BY-SA版权

文章标签：

#YOLO

于 2024-03-13 18:44:00 首次发布

本文描述了作者在将Yolov5模型从PyTorch转换到RKNN并部署过程中遇到的问题，包括代码修改、模型转换和C++中推理失败的解决方案。作者分享了从训练到部署的完整流程和关键代码调整点。

踩了三天坑，看了各种资料，跑官网提供的C++ demo，python demo可以得到正确的结果，但是从头从 pt->onnx->rknn文件总是无法成功。因为不好的工作思路，习惯找开源，于是到git 上下载各种版本代码尝试都出推理结果不正确的问题。后来头疼的想要放弃，后来看到博客给提供了思路，分别为：

rknn导出yolov5模型在python中可正常推理，在c++中无结果._rknn 模型推理无法获取结果-优快云博客 Yolov5笔记--RKNN推理部署源码的粗略理解_rknn yolov5-优快云博客

废话不多说，进入正题，给大家提供一套完整的从训练到部署的使用流程。

训练代码链接，使用官网提供的即可：

GitHub - ultralytics/yolov5: YOLOv5 🚀 in PyTorch > ONNX > CoreML > TFLite

各位需要根据实际需求进行类别修改，以及训练数据存放位置的调整，这里不再赘述。

第一步：训练好的pt 模型转换成onnx ，需要对代码作出修改，修改的位置分别包含两处：

models/yolo.py，里面Detect 类里面的forward 函数。

第一处将：

def forward(self, x):
z = [] # inference output
for i in range(self.nl):
x[i] = self.m[i](x[i]) #