YOLOV5s基于pytorch框架进行训练和部署_pytorch 环境yolov5s.pt模型的压缩

yolo_v5s训练预部署(pytorch版本)

1.下载代码和预训练模型

git clone https://github.com/ultralytics/yolov5.git
1

网址： https://github.com/ultralytics/yolov5
预训练模型用于加快收敛

2.搭建yolov5的训练环境

pip install -r requirements.txt  # install
1

或者在conda环境下新建一个虚拟环境进行相关软件的安装

# 安装pip,若虚拟环境中已经安装pip则跳过
conda install pip
# 使用conda 中的pip安装yolov5训练环境所使用的软件
anaconda3/envs/your_env/bin/pip install -r requirements.txt
1
2
3
4

参考网址： https://blog.csdn.net/dou3516/article/details/103169690

3.将数据集整理成yolov5训练需要的格式

使用脚本make_list_food_train_test_for_yolov5.py

# -*- coding: utf-8 -*-
"""
Created on Mon Jun  1 15:22:52 2020

@author: keylab
"""

#同时建立train、valid的图片路径
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) #1pingguo 2shengnvguo...
for filename in filenames: #1pingguo
    img_full_file = dataDir + filename 
    category_filenames = os.listdir(img_full_file + '/') #724.jpg 724.txt...
    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)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40

具体的会生成以下四个目录：

|————shujuji
	|————train
		|————images
		|————labels
	|————val
		|————images
		|————labels

1
2
3
4
5
6
7
8

4.设置配置文件

4.1 data/data100.yaml

在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']  # 类别名称
1
2
3
4

4.2 models/yolov5s_data100.yaml

在这个文件中只需要修改类别的数量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 ""
1
2
3
4

具体的各参数含义：
其中weights是权重文件.pt格式，可以输入空格，代表使用随机权重，或者输入权重文件的路径
cfg是模型的yaml文件，一般存放在models文件夹里
data是数据集的yaml文件，一般存放在data文件夹里
epochs是训练轮数，默认300轮
batch-size是batch数，默认16
img是输入图片大小，网络会自动按参数进行resize,默认640X640

参考链接：https://blog.csdn.net/qq_51993578/article/details/114950751

6.使用该模型进行目标检测预分类

python detect.py --weights xxx.pt --source xxx.avi --view-img
1

–view-img为显示图片

Run inference on images, videos, directories, streams, etc.

Usage - sources:
    $ python path/to/detect.py --weights yolov5s.pt --source 0              # webcam
                                                             img.jpg        # image
                                                             vid.mp4        # video
                                                             path/          # directory
                                                             path/*.jpg     # glob
                                                             'https://youtu.be/Zgi9g1ksQHc'  # YouTube
                                                             'rtsp://example.com/media.mp4'  # RTSP, RTMP, HTTP stream

Usage - formats:
    $ python path/to/detect.py --weights yolov5s.pt                 # PyTorch
                                         yolov5s.torchscript        # TorchScript
                                         yolov5s.onnx               # ONNX Runtime or OpenCV DNN with --dnn
                                         yolov5s.xml                # OpenVINO
                                         yolov5s.engine             # TensorRT
                                         yolov5s.mlmodel            # CoreML (MacOS-only)
                                         yolov5s_saved_model        # TensorFlow SavedModel
                                         yolov5s.pb                 # TensorFlow GraphDef
                                         yolov5s.tflite             # TensorFlow Lite
                                         yolov5s_edgetpu.tflite     # TensorFlow Edge TPU
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22

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)')
1
2

运行命令：

# 相关路径已经在脚本中进行设置
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
1
2
3
4

8.转换后的data100.onnx有两种方式进行部署

部署方式简介如下，具体咋搞以后有空再弄了。。。

8.1 Python环境部署

在电脑上使用pip安装onnx、onnxruntime库对模型进行部署
这一步在之前使用export.py导出onnx类型的权重时就已经进行了相关文件的安装
这里可以直接使用作者提供的detect.py文件，也可以自己进行编写

8.2 C++部署—在onnx官网下载相关安装包，进行编译安装

onnxruntime的安装包较大，可以直接将编译好的库文件和相关头文件放在运行设备上

8.3 C++部署—借助于opencv4.x中的模型推理模块实现

该方式可以在CPU上运行