目标检测：Yolov5训练自己的数据集（完整版）

2401_84819243

已于 2024-09-27 13:48:08 修改

阅读量1.1k

点赞数 15

文章标签： YOLO 计算机视觉人工智能目标检测机器学习深度学习 python

于 2024-09-27 13:08:46 首次发布

本文链接：https://blog.csdn.net/2401_84819243/article/details/140416533

版权

利用Yolov5训练自己的数据集，我们主要做以下工作：

一、源码克隆和环境依赖

（一）源码克隆

YOLOv5的代码是开源的，因此我们可以从github上克隆其源码。我们利用yolov5.5分支来作为讲解。这里不再详细介绍，将代码下载下来，至此整个项目就已经准备好了。接下来我们准备运行环境。

（二）环境依赖

本教程使用环境：

torch: 2.0.0
python: 3.8
yolov5 v6.0

这边要想成功运行YOLOv5代码，必须安装必要的依赖库文件。库文件要求，在源码中已经给出，我们按照下面操作执行就可以了。打开requirements.txt这个文件，可以看到里面有很多的依赖库和其对应的版本要求。我们打开pycharm的命令终端，在中输入如下的命令，就可以安装了。

pip install -r requirements.txt -i  https://pypi.tuna.tsinghua.edu.cn/simple

至此，深度学习的环境和依赖包就都结束了。因为本教程主详细讲解训练自己的数据集，避免篇幅过多，因此这里不详细说明了。

二、模型训练

当我们准备好必要的运行环境、代码和自己的数据集后，接下来就要替换自己的数据集，进行模型训练

（一）修改配置文件

预训练模型和数据集都准备好了，就可以开始训练自己的yolov5目标检测模型了，训练目标检测模型需要修改两个yaml文件中的参数。一个是data目录下的相应的yaml文件，一个是model目录文件下的相应的yaml文件。

修改data目录下的相应的yaml文件。找到目录下的voc.yaml文件。建议都使用绝对路径，不容易出错.
train：是上文所提到的训练train的图片
val：是上文提到的验证集valid的图片
nc：是标签的数目
name：是标签的名字，数字从0开始

（二）修改train.py文件

主要修改的：
weights：权重文件路径，看你想使用的yolov5的什么模型有s,l,n,x这样的，选其中一个就行。运行之后会自动下载，如果下载的比较慢，可以自己从网上查找
cfg：存储模型结构的配置文件，一般不用管。如果有报错说这里有错误的话，你就找你用的什么模型，在models文件夹里找到对应的.yaml文件。改一下对应的nc。就是你的标签有多少。

data：存储训练、测试数据的文件。就是上文中提到的修改的.yaml文件
epochs：指的就是训练过程中整个数据集将被迭代（训练）了多少次，显卡不行你就调小点。
workers：线程数。默认是8。

（三）修改detect.py（推理）文件

等到数据训练好了以后，就会在主目录下产生一个run文件夹，在run/train/exp/weights目录下会产生两个权重文件，一个是最后一轮的权重文件，一个是最好的权重文件，一会我们就要利用这个最好的权重文件来做推理测试。除此以外还会产生一些验证文件的图片等一些文件。

找到主目录下的detect.py文件，打开该文件。然后找到主函数的入口，这里面有模型的主要参数。模型的主要参数解析如下所示。

f __name__ == '__main__':
"""
--weights:权重的路径地址
--source:测试数据，可以是图片/视频路径，也可以是'0'(电脑自带摄像头),也可以是rtsp等视频流
--output:网络预测之后的图片/视频的保存路径
--img-size:网络输入图片大小
--conf-thres:置信度阈值
--iou-thres:做nms的iou阈值
--device:是用GPU还是CPU做推理
--view-img:是否展示预测之后的图片/视频，默认False
--save-txt:是否将预测的框坐标以txt文件形式保存，默认False
--classes:设置只保留某一部分类别，形如0或者0 2 3
--agnostic-nms:进行nms是否也去除不同类别之间的框，默认False
--augment:推理的时候进行多尺度，翻转等操作(TTA)推理
--update:如果为True，则对所有模型进行strip_optimizer操作，去除pt文件中的优化器等信息，默认为False
--project：推理的结果保存在runs/detect目录下
--name：结果保存的文件夹名称
"""
    parser = argparse.ArgumentParser()
    parser.add_argument('--weights', nargs='+', type=str, default='yolov5s.pt', help='model.pt path(s)')
    parser.add_argument('--source', type=str, default='data/images', help='source')  # file/folder, 0 for webcam
    parser.add_argument('--img-size', type=int, default=640, help='inference size (pixels)')
    parser.add_argument('--conf-thres', type=float, default=0.25, help='object confidence threshold')
    parser.add_argument('--iou-thres', type=float, default=0.45, help='IOU threshold for NMS')
    parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
    parser.add_argument('--view-img', action='store_true', help='display results')
    parser.add_argument('--save-txt', action='store_true', help='save results to *.txt')
    parser.add_argument('--save-conf', action='store_true', help='save confidences in --save-txt labels')
    parser.add_argument('--nosave', action='store_true', help='do not save images/videos')
    parser.add_argument('--classes', nargs='+', type=int, help='filter by class: --class 0, or --class 0 2 3')
    parser.add_argument('--agnostic-nms', action='store_true', help='class-agnostic NMS')
    parser.add_argument('--augment', action='store_true', help='augmented inference')
    parser.add_argument('--update', action='store_true', help='update all models')
    parser.add_argument('--project', default='runs/detect', help='save results to project/name')
    parser.add_argument('--name', default='exp', help='save results to project/name')
    parser.add_argument('--exist-ok', action='store_true', help='existing project/name ok, do not increment')
    opt = parser.parse_args()

这里需要将刚刚训练好的最好的权重传入到推理函数中去。然后就可以对图像视频进行推理了。

parser.add_argument('--weights', nargs='+', type=str, default='runs/train/exp/weights/best.pt', help='model.pt path(s)')

对图片进行测试推理，将如下参数修改成图片的路径，然后运行detect.py就可以进行测试了。

 parser.add_argument('--source', type=str, default='000295.jpg', help='source')

推理测试结束以后，在run下面会生成一个detect目录，推理结果会保存在exp目录下。如图所示。

图片的推理结果如下所示。效果还是很不错的。

至此yolov5训练自己的模型就完全搞定了。

三、训练结果讲解

模型训练结束以后，在run下面会生成一个train目录，推理结果会保存在exp目录下。如图所示。

从上述图片中可以看到的信息是：
我总共训练了200次，共花费3.739个小时。
最新的训练结果是：last.pt，文件位置如图所示，文件大小是27.4MB
最好的训练结果是：best.pt，文件位置如图所示，文件大小是27.4MB
结果保存在runs/train/exp14
打开该文件，如下图所示

（一）深入了解文件含义：损失函数、mAP可视化结果

这个图片内有YOLOv5的损失，其中损失分为三种

定位损失box loss: YOLO V5使用 GIOU Loss作为bounding box的损失，Box推测为GIoU损失函数均值，box loss的值越小框的位置越准
置信度损失obj loss: 推测为目标检测loss均值，值越小目标检测的越准
分类损失cls loss: 推测为分类loss均值，值越小预测的类别越准
val/boxloss: 验证集bounding box损失
val/obj loss: 验证集目标检测loss均值
val/cls loss: 验证集分类loss均值，我训练时只有person这一类，所以为0【单类别这个值都是0】
mAP@0.5:0.95: 表示在不同的10U阈值(从0.5到0.95，步长为0.05)(0.5、0.55、0.6、0.65、0.7、0.75、0.8、0.85、0.9、0.95)上的mAP
mAP@0.5: 表示阈值大于0.5的平均mAP