1.下载代码和预训练模型
git clone https://github.com/ultralytics/yolov5.git
2.搭建yolov5的训练环境
pip install -r requirements.txt
或者在conda环境下新建一个虚拟环境进行相关软件的安装
3.将数据集整理成yolov5训练需要的格式
使用脚本 make_list_food_train_test_for_yolov5.py
"""
Created on Mon Jun 1 15:22:52 2020
@author: keylab
"""
import os
import os.path
import shutil
dataDir='your_path'
train_image_dir = 'your_path'
train_label_dir = 'your_path'
val_image_dir = 'your_path'
val_label_dir = 'your_path'
filenames = os.listdir(dataDir)
for filename in filenames:
img_full_file = dataDir + filename
category_filenames = os.listdir(img_full_file + '/')
count = 11
for category_filename in category_filenames:
category_tem = category_filename[:-4]
if category_filename[-3:]=='txt':
src_txt = img_full_file + '/' + category_tem + '.txt'
src_img = img_full_file + '/' + category_tem + '.jpg'
count=count+1
if count %10 == 0:
dst_txt_val = val_label_dir + category_tem + '.txt'
dst_img_val = val_image_dir + category_tem + '.jpg'
shutil.copyfile(src_txt, dst_txt_val)
shutil.copyfile(src_img, dst_img_val)
else:
dst_txt_train = train_label_dir + category_tem + '.txt'
dst_img_train = train_image_dir + category_tem + '.jpg'
shutil.copyfile(src_txt, dst_txt_train)
shutil.copyfile(src_img, dst_img_train)
具体的会生成以下四个目录:
|————shujuji
|————train
|————images
|————labels
|————val
|————images
|————labels
4.设置配置文件
在data目录下新建一个data.yaml文件,文件名根据自己需要进行自定义;
该文件只需要保留4行内容,分别为:
train: your_path/train/images
val: your_path/val/images
nc: 11
names: ['kong','pingguo','shengnvguo','qiezi','xihongshi','lvyuanjiao','huanggua','beibeinangua','youcai','huluobo','wawacai']
在这个文件中只需要修改类别的数量 nc
,这里改为11
5.开始训练
python train.py --img 640 --batch 4 --epochs 300 --data data/data100.yaml --cfg models/yolov5s_data100.yaml --weights weigth/yolov5s.pt
python train.py --img 640 --batch 4 --epochs 300 --data data/data100.yaml --cfg models/data100.yaml --weights ""
具体的各参数含义:
其中 weights
是权重文件 .pt
格式,可以输入空格,代表使用随机权重,或者输入权重文件的路径
cfg
是模型的 yaml
文件,一般存放在 models
文件夹里
data
是数据集的 yaml
文件,一般存放在 data
文件夹里
epochs
是训练轮数,默认 300
轮
batch-size
是 batch
数,默认 16
img
是输入图片大小,网络会自动按参数进行resize,默认 640X640
6.使用该模型进行目标检测预分类
python detect.py --weights xxx.pt --source xxx.avi --view-img
–view-img为显示图片
Run inference on images, videos, directories, streams, etc.
Usage - sources:
$ python path/to/detect.py --weights yolov5s.pt --source 0
img.jpg
vid.mp4
path/
path/*.jpg
'https://youtu.be/Zgi9g1ksQHc'
'rtsp://example.com/media.mp4'
Usage - formats:
$ python path/to/detect.py --weights yolov5s.pt
yolov5s.torchscript
yolov5s.onnx
yolov5s.xml
yolov5s.engine
yolov5s.mlmodel
yolov5s_saved_model
yolov5s.pb
yolov5s.tflite
yolov5s_edgetpu.tflite
7.将训练好的data100.pt转换为data100.onnx
yolovs自带了onnx模型转换脚本,即 export.py
打开该脚本,在 parse_opt()
函数相应的yaml和权重路径处进行相应的修改,也可以在终端自行输入
parser.add_argument('--data', type=str, default=ROOT / 'data/data100.yaml', help='dataset.yaml path')
parser.add_argument('--weights', nargs='+', type=str, default=ROOT / 'backup/yolov5s_data100.pt', help='model.pt path(s)')
运行命令:
python3 export.py --img-size 640 --batch-size 1
python3 export.py --weights backup/yolov5s_data100.pt data/data100.yaml --img-size 640 --batch-size 1
8.转换后的data100.onnx有两种方式进行部署
部署方式简介如下,具体咋搞以后有空再弄了。。。
在电脑上使用pip安装onnx、onnxruntime库对模型进行部署
这一步在之前使用 export.py
导出onnx类型的权重时就已经进行了相关文件的安装
这里可以直接使用作者提供的 detect.py
文件,也可以自己进行编写
onnxruntime的安装包较大,可以直接将编译好的库文件和相关头文件放在运行设备上
该方式可以在CPU上运行
Original: https://blog.csdn.net/guanjing_dream/article/details/123715084
Author: 我是快乐的小趴菜
Title: YOLOV5s基于pytorch框架进行训练和部署
原创文章受到原创版权保护。转载请注明出处:https://www.johngo689.com/687541/
转载文章受原作者版权保护。转载请注明原作者出处!