YOLOV5s基于pytorch框架进行训练和部署

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-sizebatch数,默认 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/

转载文章受原作者版权保护。转载请注明原作者出处!

(0)

大家都在看

亲爱的 Coder【最近整理,可免费获取】👉 最新必读书单  | 👏 面试题下载  | 🌎 免费的AI知识星球