用YOLOv5代码训练模型并部署到Android客户端


仅做记录。


YOLOv5代码来源

训练模型时遇到的问题

(1)RuntimeError: CUDA error: no kernel image is available for execution on the device

报错信息显示显卡的CUDA计算能力和pytorch版本不匹配
安装适用于使用显卡的pytorch
卸载当前版本的pytorch, 重新安装匹配版本

pip uninstall torch
pip uninstall torchvision
pip install torch==1.7.0+cu110 torchvision==0.8.1+cu110 torchaudio===0.7.0 -f https://download.pytorch.org/whl/torch_stable.html -i  https://pypi.tuna.tsinghua.edu.cn/simple/

注意上面的内容要与自己的cuda版本匹配哈,怎么查看当前环境的cuda版本,命令如下:

查看可用cuda最大可用版本:

nvidia-smi

在这里插入图片描述
查看当前环境的cuda版本:

nvcc -V

在这里插入图片描述

(2)在vscode上使用tensorboard

(1)找到log文件的保存地址:
在这里插入图片描述
(2)vscode连接上服务器之后,在终端输入:

tensorboard --logdir=/databank/home/**/**/yolov5/runs/train/exp12/   #这里的地址为上述地址

在这里插入图片描述
这样就可以啦!!!

将YOLOv5训练得到的模型部署到Android客户端:

配置好环境,使用如下命令,运行train.py文件得到“.pt”模型

YOLOv5代码来源

CUDA_VISIBLE_DEVICES=0 python train.py --img 640 --batch 16 --epochs 300 --data  data/my.yaml  --weights weights/yolov5s.pt

在这里插入图片描述

使用export.py中代码将“.pt”模型转为“.torchscript.ptl”模型

注意,原export.py中只能将“.pt”模型转为“.torchscript.pt”,不能转为“.torchscript.ptl”,所以需要稍微修改下代码,具体如下:

在这里插入图片描述
把蓝色框中代码注释掉,并写入红色框中代码:

f = file.with_suffix('.torchscript.ptl')
(optimize_for_mobile(ts) if optimize else ts)._save_for_lite_interpreter(str(f))

然后使用如下命令:

python export.py --weights runs/train/exp11/weights/best.pt --include torchscript

在这里插入图片描述
在这里插入图片描述

将生成的best.torchscript.ptl模型部署到android客户端进行目标检测

Git中提供的android-demo-app模板代码
(1)在MainActivity.java文件中改变自己的图片列表:
在这里插入图片描述

private String[] mTestImages = {"aicook1.jpg", "aicook2.jpg", "aicook3.jpg", "test1.png", "test2.jpg", "test3.png"};

(2)在MainActivity.java文件中改成自己的模型:
在这里插入图片描述

mModule = LiteModuleLoader.load(MainActivity.assetFilePath(getApplicationContext(), "best.torchscript.ptl"));
BufferedReader br = new BufferedReader(new InputStreamReader(getAssets().open("aicook.txt")));

(3)将PrePostProcessor.java中的mOutputColumn的数值改为自己模型的预测种类+5,例如,如果自己模型可以检测的目标种类为30,则值为30+5=35;
在这里插入图片描述

private static int mOutputColumn = 35;

(4)将ObjectDetectionActivity.java文件中的模型也给成自己的:
在这里插入图片描述

mModule = LiteModuleLoader.load(MainActivity.assetFilePath(getApplicationContext(), "yolov5s.torchscript.ptl"));

(5)运行demo_app
在这里插入图片描述


评论 6
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值