前言
一、PT-->ONNX-->NCNN
1、准备pt文件
2、PT转ONNX
3、简化ONNX
4、ONNX转param、bin
5、修改param
二、Android 平台部署
1、安装Android studio
2、下载ncnn-android-yolov5案例
3、修改example以适应自己的数据模型
总结
参考
接上篇YOLOv5学习记录(一): 训练自己的数据模型
在上一篇文章中,我们大致讲述了如何在YOLO框架下训练自己的数据模型。
本篇内容进一步讲述如何将训练得到的数据模型,转换格式并部署在Android平台上。
按照上篇文章的步骤,想必大家都已经得到了自己训练好的模型,我的训练好的模型在./runs/train/exp4/weights下名为best.pt的文件;
在yolov5-master根目录下有export.py文件,可以直接采用转换命令:
python export.py --data ./data/myeye.yaml --weights ./runs/train/exp4/weights/best.pt --img 640 --batch 1 --train
注意:转换命令后加:--train;(去掉后处理部分的导出,不然后面转param会报错)
YOLOv5 6.0 去掉了focus,不需要手工修改任何文件,在ncnn使用时兼容性更好。网上很多教程都是基于6.0之前的版本,转换命令和后面修改./models/common.py办法不能照搬(这个坑困扰了很久emo)
最终在大佬的文章
详细记录u版YOLOv5目标检测ncnn实现评论区看到了一句话:
运行成功,会在与best.pt同目录下产生两个文件best.onnx和best.torchscript(也是一种模型转换格式),我们这里只讨论onnx。
此时可以通过可视化工具Netron查看网络的总体架构,如果不想安装到本地,可以直接点击Open Model,选择刚才转换出来的best.onnx文件即可。
安装命令如下:
pip install netron -i https://pypi.tuna.tsinghua.edu.cn/simple
使用方法如下:
可视化网络结构:
转为onnx格式后,还需要再进行模型简化,不然会在转换的时候报错:
安装简化工具:
pip install onnx-simplifier
简化命令如下:
python -m onnxsim best.onnx best-sim.onnx
此时就已经基本完成了。
其实不用按照网上的有些教程下载ncnn源码编译配置,说实话,这个编译配置过程真的很容易让人劝退。(我就转个格式,给我搞这么多花里胡哨?)
这里推荐两种比较简便的转换方式:
1、下载编译好的:Releases · Tencent/ncnn · GitHub,可以根据自己的平台和配置选择合适的编译好的文件。以ncnn-VS2019为例,可以直接使用在路径./x64/bin下的onnx2ncnn.exe来实现转换,转换命令如下(记得修改自己的路径):
onnx2ncnn.exe best-sim.onnx best.param best.bin
转换会生成两个文件,param是网络结构文件,bin为网络权重文件。
2、在查找资料的过程中也发现了一个在线转换神器(不知道会不会收集自己的模型,但是确实好用):
一键转换 Caffe, ONNX, TensorFlow 到 NCNN, MNN, Tengine (convertmodel.com)
如果出现转换报错,请再次确认模型是否经过简化、导出命令是否加入:--train
打开刚才生成的.param文件,拉到最后,将红框内的三个数据全部修改为:-1,这样做是为了防止移植到Android端,检测时出现密密麻麻的检测框覆盖图片。
至此,模型转换完毕,让我们一起去部署吧!
自己训练出来的模型,当然还是要自己来移植部署,所以又解锁了新技能:Android studio;
安装下载请参考:Android Studio安装教程(超级详细)_dhhyx的博客-CSDN博客_android studio安装教程
NCNN官方提供的example:GitHub - nihui/ncnn-android-yolov5: The YOLOv5 object detection android example
下载到本地后,还需要按照步骤进行配置:
(注:安装配置时注意待部署手机支持的SDK最低版本,此处如果出现报错,大概率是因为版本匹配问题,确定版本号后,修改build.gradle中的defaultConfig的minSdkVersion等参数即可)
1、anchors的值在训练开始时是默认根据数据集重新计算,所以我们首先要拿到数据模型的anchors值,代码内容如下:
import torch
from models.experimental import attempt_load
model = attempt_load('./weights/best.pt', map_location=torch.device('cpu'))
m = model.module.model[-1] if hasattr(model, 'module') else model.model[-1]
print(m.anchor_grid)
输出如下:
tensor([[[[[[ 6.75781, 2.00000]]],
[[[ 8.60938, 2.83594]]],
[[[ 11.01562, 4.14844]]]]],
[[[[[ 20.28125, 5.57812]]],
[[[ 26.29688, 9.30469]]],
[[[ 37.43750, 8.21875]]]]],
[[[[[39.31250, 12.95312]]],
[[[54.50000, 16.57812]]],
[[[68.62500, 25.54688]]]]]])
当然,你也可以选择在训练开始的时候关闭自适应计算训练集中的最佳锚框值。
将train.py中上面的一行代码,设置为False即可,每次训练时就不会自动计算。(这个我没有试过)
2、同时输出层ID在Permute层里面,找到对应的参数:
最后,将模型对应的anchors值、及Permute层的参数改写到ncnn-android-yolov5-master/app/src/main/jni/yolov5ncnn_jni.cpp文件中,如下图所示:
至此。所有工作准备完毕,可以在Android上测试自己的模型啦!
以上就是模型部署的全部内容,各位有什么问题欢迎大家在评论区留言讨论,互相学习交流!
1、目标检测 YOLOv5 anchor设置_深度学习
2、详细记录u版YOLOv5目标检测ncnn实现
3、Windows环境下yolov5移植到Android手机上的详细步骤