利用yolov5训练自定义数据集

目录

一、前言

二、yolov5训练的大致流程

三、具体步骤

1. 制作YOLOv5格式数据集

A. 准备数据和标签

B. 制作数据配置文件

2. 修改文件

3. 训练

4. 测试 

四、总结

---

一、前言

        随着基于深度学习算法的不断发展，yolo系列算法得到了广泛的研究。本文介绍了yolov5算法如何训练自己的数据集，并成功应用于自己的项目当中，解决复杂场景中的目标检测问题。

二、yolov5训练的大致流程

1.准备好项目所需的数据和标签（数据和标签一一对应)；

2.数据配置文件（data.yaml）需配置好训练集和验证集的路径，修改类别的数量和名称；

3.在yolov5源码中打开，train.py文件进行配置：

代码如下：

--data ../orange_detection1/data.yaml #数据配置文件
--cfg models/yolov5s.yaml #网络模型
--batch-size 16 #批处理数量
--epochs 300 #训练次数

4.进行训练，得到best.pt和last.pt模型权重。

三、具体步骤

1. 制作YOLOv5格式数据集

A. 准备数据和标签

（1）首先需要在自己电脑（linux或win10）安装LabelImg，用于人工标注数据制作数据集。

打标数据并制作数据集可参考我的这篇博客： 

 LabelImg-制作YOLOv5格式的数据集（win10）_现实话费的博客-CSDN博客_labelimg yolov5

（2）建立自定义数据集文件夹，里面分别建立test、train、valid三个文件夹，并在其每个文件夹下建立images和labels文件夹，保存图片和对应的标签文件，本文我建的数据集（orange-detection1）架构如下所示。并把自定义的数据集文件夹放入到yolov5同级目录下。

B. 制作数据配置文件

         制作data.yaml文件（注意训练和验证图片的路径），并放在数据集（orange-detection1）文件夹下。data.yaml内容如下：

train: ../orange-detection1/train/images #训练集路径
val: ../orange-detection1/valid/images  #验证集路径

nc: 1                                    #目标类别
names: ['orange']                        #目标名称

2. 修改文件

（1）修改模型配置文件yolov5s.yaml里的nc（目标的类别数量）

由于我的数据集只有一类目标，所以设置 nc: 1。

（2）修改train.py（可不修改，直接通过命令来更改参数进行训练）

在yolov5/train.py里修改如下参数：（batch-size的大小根据自己电脑显卡而定）

--data ../orange_detection1/data.yaml    #数据配置文件
--cfg models/yolov5s.yaml     #网络模型
--batch-size 16               #批处理数量
--epochs 300                  #训练次数

train.py里待修改的部分：

（3）将下载好的预训练权重文件（.pt）放入weights文件夹下。

3. 训练

（1）进入虚拟环境

conda activate yolov5    # yolov5为创建的虚拟环境的名字

（2）进入项目路径

cd work1/yolov5

（3）下载项目所需要的库

首先注释掉 torch、torchvision，因为配置环境时，我们下载好了pytorch GPU环境。

# torch>=1.6.0
# torchvision>=0.7.0

然后运行命令：pip install -U -r requirements.txt    #下载yolov5需要的库（如果执行该命令，下载过程中报错了，那么就请手动安装吧）。

yolov5源码下有requirements.txt文件，在文件夹下查找，并手动安装：pip install xxxxx。

requirements.txt文件如下：

（4）训练模型（以yolov5s为例）

如果没有修改train.py，可以在命令中直接添加参数指令来进行训练；如果修改了train.py，则直接运行python train.py指令来训练模型。

A. 基于训练得到的best.pt模型来训练：

$ python train.py --img 640 --batch 16 --epochs 300 --data ../orange-detection1/data.yaml --cfg models/yolov5s.yaml --weights best.pt

 B. 基于官网的预训练模型来训练：

$ python train.py --img 640 --batch 16 --epochs 300 --data ../orange-detection1/data.yaml --cfg models/yolov5s.yaml --weights yolov5s.pt

 C. 不使用预训练权重，从头开始训练：

$ python train.py --img 640 --batch 16 --epochs 300 --data ../orange-detection1/data.yaml --cfg models/yolov5s.yaml --weights ''

训练过程： 

4. 测试 

        训练结束后，在yolov5s/runs/train/exp/weights文件夹下生成了best.pt和last.pt，利用自定义数据集中的测试集图片（orange-detection1/test/images/图片名）来测试模型的效果。

检测一张图片： 

python detect.py --weight weights/best.pt --source ../orange-detection1/test/images/图片名

批量检测测试集的全部图片：

python detect.py --weight weights/best.pt --source ../orange-detection1/test/images

用训练得到的最好的best.pt模型检测测试集下模型mAP：

python test.py --data ../orange_detection1/data.yaml --weight weights/best.pt --augment

        执行成功后，会显示检测结果的保存路径，并在保存路径（yolov5s/runs/detect(或test)）下查看检测图片和测试结果。

---

四、总结

        总的来说，yolov5算法的训练相对简单，希望各位小伙伴成功训练出性能较好的模型，欢迎大家和我交流讨论。