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=int, default=640, help='train, val image size (pixels)')
    parser.add_argument('--rect', action='store_true', help='rectangular training')
    parser.add_argument('--resume', nargs='?', const=True, default=False, help='resume most recent training')
    parser.add_argument('--nosave', action='store_true', help='only save final checkpoint')
    parser.add_argument('--noval', action='store_true', help='only validate final epoch')
    # 自动计算锚框，设置成False，则不会自动计算
    parser.add_argument('--noautoanchor', action='store_true', help='disable autoanchor check')
    parser.add_argument('--evolve', type=int, nargs='?', const=300, help='evolve hyperparameters for x generations')
    parser.add_argument('--bucket', type=str, default='', help='gsutil bucket')
    parser.add_argument('--cache', type=str, nargs='?', const='ram', help='--cache images in "ram" (default) or "disk"')
    parser.add_argument('--image-weights', action='store_true', help='use weighted image selection for training')
    # 是否用GPU
    parser.add_argument('--device', default='0', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
    parser.add_argument('--multi-scale', action='store_true', help='vary img-size +/- 50%%')
    parser.add_argument('--single-cls', action='store_true', help='train multi-class data as single-class')
    parser.add_argument('--adam', action='store_true', help='use torch.optim.Adam() optimizer')
    parser.add_argument('--sync-bn', action='store_true', help='use SyncBatchNorm, only available in DDP mode')
    parser.add_argument('--workers', type=int, default=0, help='maximum number of dataloader workers')
    parser.add_argument('--project', default='runs/train', help='save to project/name')
    parser.add_argument('--entity', default=None, help='W&B entity')
    parser.add_argument('--name', default='exp', help='save to project/name')
    parser.add_argument('--exist-ok', action='store_true', help='existing project/name ok, do not increment')
    parser.add_argument('--quad', action='store_true', help='quad dataloader')
    parser.add_argument('--linear-lr', action='store_true', help='linear LR')
    parser.add_argument('--label-smoothing', type=float, default=0.0, help='Label smoothing epsilon')
    parser.add_argument('--upload_dataset', action='store_true', help='Upload dataset as W&B artifact table')
    parser.add_argument('--bbox_interval', type=int, default=-1, help='Set bounding-box image logging interval for W&B')
    parser.add_argument('--save_period', type=int, default=-1, help='Log model after every "save_period" epoch')
    parser.add_argument('--artifact_alias', type=str, default="latest", help='version of dataset artifact to be used')
    parser.add_argument('--local_rank', type=int, default=-1, help='DDP parameter, do not modify')
    parser.add_argument('--freeze', type=int, default=0, help='Number of layers to freeze. backbone=10, all=24')
    parser.add_argument('--patience', type=int, default=300, help='EarlyStopping patience (epochs without improvement)')
    opt = parser.parse_known_args()[0] if known else parser.parse_args()
    return opt

主要是修改 --weights, --cfg, --data, --epochs 和 --batch-size
修改完就可以train啦
train完生成 run/train 文件夹，里面保存有训练好的权重，best.pt 和 last.pt，按我的经验，best.pt的测试效果更好一些，可以都试试。

4. 测试–val.py

这里注意，因为我们的测试文件是放在test.txt里面的，所以在val.py里面需要修改102行、156行、166行，val都改成test：

task='test',  # train, val, test, speed or study

is_coco = isinstance(data.get('test'), str) and data['test'].endswith('coco/val2017.txt')  # COCO dataset

task = task if task in ('train', 'val', 'test') else 'test'  # path to train/val/test images

如果你的测试集是放在val.txt里面就不用修改了
然后改配置：

def parse_opt():
    parser = argparse.ArgumentParser()
    # 修改成之前修改的数据配置文件
    parser.add_argument('--data', type=str, default='data/mydata.yaml', help='dataset.yaml path')
    # 修改成run/train里面生成的权重文件
    parser.add_argument('--weights', nargs='+', type=str, default='run/train/exp/last.pt', help='model.pt path(s)')
    parser.add_argument('--batch-size', type=int, default=32, help='batch size')
    parser.add_argument('--imgsz', '--img', '--img-size', type=int, default=640, help='inference size (pixels)')
    parser.add_argument('--conf-thres', type=float, default=0.001, help='confidence threshold')
    parser.add_argument('--iou-thres', type=float, default=0.6, help='NMS IoU threshold')
    # 因为测试集是在test.txt里面，所以val改成test
    parser.add_argument('--task', default='test', help='train, val, test, speed or study')
    parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
    parser.add_argument('--workers', type=int, default=8, help='max dataloader workers (per RANK in DDP mode)')
    parser.add_argument('--single-cls', action='store_true', help='treat as single-class dataset')
    parser.add_argument('--augment', action='store_true', help='augmented inference')
    parser.add_argument('--verbose', action='store_true', help='report mAP by class')
    parser.add_argument('--save-txt', action='store_true', help='save results to *.txt')
    parser.add_argument('--save-hybrid', action='store_true', help='save label+prediction hybrid results to *.txt')
    parser.add_argument('--save-conf', action='store_true', help='save confidences in --save-txt labels')
    parser.add_argument('--save-json', action='store_true', help='save a COCO-JSON results file')
    parser.add_argument('--project', default=ROOT / 'runs/val', help='save to project/name')
    parser.add_argument('--name', default='exp', help='save to project/name')
    parser.add_argument('--exist-ok', action='store_true', help='existing project/name ok, do not increment')
    parser.add_argument('--half', action='store_true', help='use FP16 half-precision inference')
    parser.add_argument('--dnn', action='store_true', help='use OpenCV DNN for ONNX inference')

主要是改 --data, --weights, --task这三个，如果你的测试集比较少也可以把 --batch-size 调小一些。

5. 检测–detect.py

修改下面三个就可以了

# 改成训练好的权重文件
parser.add_argument('--weights', nargs='+', type=str, default='run/train/exp/last.pt', help='model path(s)')
#改成要检测的图片路径
parser.add_argument('--source', type=str, default='data/images', help='file/dir/URL/glob, 0 for webcam')
# 改成数据配置文件
parser.add_argument('--data', type=str, default='data/mydata.yaml', help='(optional) dataset.yaml path')

总结

如果是要改进yolov5网络的话，就要修改model文件夹的.yaml文件，希望大家都可以把yolov5网络用懂、看懂~

基本上就是这么多，如果还有什么问题欢迎在下面留言，我看到就会回复哒~