本专栏主要是提供一种国产化图像识别的解决方案,专栏中实现了YOLOv5/v8在国产化芯片上的使用部署,并可以实现网页端实时查看。根据自己的具体需求可以直接产品化部署使用。
B站配套视频:https://www.bilibili.com/video/BV1or421T74f
基础背景
我们在之前的课程里面已经把如何训练和如何部署运行说了一下。接下来就是将这两块内容进行拼接,如何使用自己的数据集进行训练导出可以在RK3588上运行的模型。
还不知道如何训练的朋友可以看我之前的CSDN博客,或者看的B站视频。
在国产芯片上实现YOLOv5/v8图像AI识别-【1.3】YOLOv5的介绍及使用(训练、导出)
在国产芯片上实现YOLOv5/v8图像AI识别-【2.2】RK3588上C++开发环境准备及测试
PS:如果想要基础代码的可以私信我,我会逐一发送给大家。
模型修改和导出onnx
在之前的课程里面说过整个推理过程分为预处理、推理、后处理三个部分。对于常规使用GPU的yolo项目来说,可以把后处理放在GPU上运行,所以模型部分会有后处理部分的逻辑描述。而NPU受到性能限制,所以需要将后处理进行剥离。
所以我们需要对后处理部分模型结构进行修改。此处给出的是YOLOv5版本的模型修改方式,关于V8的修改方式会在后面课程里面说明。
知道原理后我们大致可以知道,需要修改后处理部分代码,和导出部分的代码。我们首先修改后处理部分代码。…/yolov5/models/yolo.py大概在50~60行左右的位置。
在forward中直接将推理的结果信息全部输出,不需要模型继续操作。
z = [] # inference output
for i in range(self.nl):
z.append(self.m[i](x[i]))
return z
之后修改根目录下的export.py,由于我们将后处理部分的结构进行了调整,变成了一个队列,所以我们需要将队列可以输出。
完成以上操作之后就可以按照之前的课程内容进行模型导出,导出成onnx模型。
python export.py --weights yolov5s.pt --include onnx --opset 11
重点提醒:在进行再次训练时候一定要还原yolo.py的修改内容否则训练会报错。
通过这种方法,我们就可以看到导出的模型产生了变化。
修改后的onnx模型修改到此已经完成,在下一篇博客我们会将模型从onnx量化成rknn。