yolov5-pytorch训练自己的数据集-超详细

---

前言

代码地址：https://github.com/ultralytics/yolov5

用yolov5有一段时间了，对这部分代码也比较熟悉一些了，写篇文章记录一下吧~

本文主要是介绍怎么利用官方代码训练与测试自己的数据集，对原理和代码的详细介绍可能少一些，因为已经有很多大牛们介绍过原理方面的知识了。

yolov5主要是有四种模型，分别是yolov5s，yolov5m，yolov5l，yolov5x，就跟我们买衣服的码一样的标号哈哈，网络的宽度深度也是依次递增。

---

一、主要代码介绍

1.data文件夹

data文件夹里面主要是放我们的图片、标注文件(.xml)和数据说配置文件 — 即指明数据存放位置(.yaml)，当然你也可以自定义一个放数据集的文件，我是自定义了一个mydata文件夹专门放数据集，不然都放data文件夹有点乱。

2.model文件夹

model文件夹里主要放的是四种网络(s, m, l, x)的配置文件(.yaml)；
common.py是每个模块的代码，如果要改进网络结构，添加新的模块可以模仿里面的代码写；
yolo.py是用于查看模型的结构和参数数量的，改一下299行的cfg，如改成下面的就可以查看yolov5x的模型结构和参数数量了。

parser.add_argument('--cfg', type=str, default='yolov5x.yaml', help='model.yaml')

3.主要py文件

1. train.py用于训练模型
2. val.py用于测试结果，输出precision, recall, mAP等信息
3. detect.py用于对图片进行检测

4.新建weights文件夹

需要新建weights文件夹，放预训练的权重文件
下载地址：https://github.com/ultralytics/yolov5/releases

训练yolov5s模型则下载yolov5s.pt权重

二、数据集准备

我是新建了一个mydata文件夹专门存放数据的，.xml文件我是用labelImg进行标注的，网上有很多资料，我这里就默认大家都会用了(偷懒.jpg)，数据集我是按下面方式存放的，当然其它方式也可以，改一下相关路径就行了。

-mydata(主目录)
–annotations(一级子目录)
–img_lab(一级子目录)
—images(二级子目录)
(说明：images文件夹放在img_lab文件夹里面，存放图片，annotations存放标注文件(.xml文件))

另外有两个.py文件用于生成labels和训练、验证、测试的txt说明文件，我放到另一篇文章里了

https://blog.csdn.net/weixin_50113231/article/details/123571036?spm=1001.2014.3001.5502
(训练yolo的数据集准备)

三、训练与测试

1. 训练前准备

首先安装各种库和包，在cmd或者python的终端(Terminal)运行：

pip install -r requirements.txt

注意要切换到yolov5的目录下

2. 修改配置文件

修改data文件夹里面的.yaml文件，可以新建一个也可以直接修改原来有的，我是新建了一个 mydata.yaml 文件：

# train and val data as 1) directory: path/images/, 2) file: path/images.txt, or 3) list: [path1/images/, path2/images/]
train: ./mydata/img_lab/train.txt
val: ./mydata/img_lab/val.txt
test: ./mydata/img_lab/test.txt
nc: 1  # 修改为自己类别数

# class names
names: ['person']

3. 训练–train.py

修改train.py里面的设置，即parse_opt

def parse_opt(known=False):
    parser = argparse.ArgumentParser()
    # 修改成模型对应的权重文件
    parser.add_argument('--weights', type=str, default='weights/yolov5s.pt', help='initial weights path')
    # 修改成要训练的模型的配置文件
    parser.add_argument('--cfg', type=str, default='models/yolov5s.yaml', help='model.yaml path')
    # 上一节里修改的数据配置文件
    parser.add_argument('--data', type=str, default='data/mydata.yaml', help='dataset.yaml path')
    # 超参数设置，不需要修改
    parser.add_argument('--hyp', type=str, default='data/hyps/hyp.scratch.yaml', help='hyperparameters path')
    # 训练多少代
    parser.add_argument('--epochs', type=int, default=500)
    # 每批多少张图片
    parser.add_argument('--batch-size', type=int, default=64, help='total batch size for all GPUs')
    parser.add_argument('--imgsz', '--img', '--img-size', type