在自定义的数据集上训练YOLOv5详细教程分享

AI算法与图像处理

共 5668字,需浏览 12分钟

 · 2020-09-03



点击上方AI算法与图像处理”,选择加"星标"或“置顶”

重磅干货,第一时间送达

来源 | 深度学习与计算机视觉

编辑 | 磐怼怼

YOLO系列的目标检测模型随着YOLOv5的引入变得越来越强大了。在这篇文章中,我们将介绍如何训练YOLOv5为你识别自己定制的对象。

本文我们使用公共血细胞检测数据集,你可以自己导出,也可以在自己自定义数据上使用本教程。
  • 公共血细胞检测数据集:https://public.roboflow.ai/object-detection/bccd
为了训练检测器,我们采取以下步骤:
  • 安装YOLOv5依赖项
  • 下载自定义YOLOv5对象检测数据
  • 定义YOLOv5模型配置和架构
  • 训练一个定制的YOLOv5检测器
  • 评估YOLOv5性能
  • 可视化YOLOv5训练数据
  • 对测试图像使用YOLOv5进行推断
  • 导出并保存YOLOv5权重以供将来使用

YOLOv5的新技术点

就在两个月前,我们对googlebrain引入EfficientDet感到非常兴奋,并写了一些关于EfficientDet的博客文章。我们认为这个模型可能会超越YOLO家族在实时目标探测领域的突出地位,但事实证明我们错了。
三周内,YOLOv4在Darknet框架下发布,我们还写了很多关于YOLOv4技术解析的文章。
在写这些文章的几个小时之前,YOLOv5发布了,我们发现它非常清晰明了。
YOLOv5是在Ultralytics-Pythorch框架中编写的,使用起来非常直观,推理速度非常快。事实上,我们和许多人经常将YOLOv3和YOLOv4 Darknet权重转换为Ultralytics PyTorch权重,以便使用更轻的库来更快地进行推理。
YOLOv5比YOLOv4表现更好吗?我们很快会向你介绍,在此之前你需要已经对YOLOv5和YOLOv4有了初步的了解。

YOLOv5与EfficientDet的性能对比
YOLOv4显然没有在YOLOv5存储库中进行评估,也就是说,YOLOv5更易于使用,而且它在我们最初运行的定制数据上表现非常出色。
我们建议你在 YOLOv5 Colab Notebook 中同时进行接下来的操作。
  • https://colab.research.google.com/drive/1gDZ2xcTOgR39tGGs-EZ6i3RTs16wmzZQ


安装YOLOv5环境

首先我们克隆YOLOv5存储库并安装依赖项,这会设置我们的编程环境,准备好运行对象检测训练和推理命令。
!git clone https://github.com/ultralytics/yolov5 # clone repo!pip install -U -r yolov5/requirements.txt # install dependencies
%cd /content/yolov5
然后,我们可以看看谷歌Colab免费提供给我们的训练环境。
import torchfrom IPython.display import Image # for displaying imagesfrom utils.google_utils import gdrive_download # for downloading models/datasets
print('torch %s %s' % (torch.__version__, torch.cuda.get_device_properties(0) if torch.cuda.is_available() else 'CPU'))
你会从谷歌Colab收到一个 Tesla P100 GPU。以下是我收到的:
torch 1.5.0+cu101 _CudaDeviceProperties(name='Tesla P100-PCIE-16GB', major=6, minor=0, total_memory=16280MB, multi_processor_count=56)
GPU可以让我们加快训练时间,Colab预装了torchcuda。如果你尝试在本地使用本教程,可能需要执行其他步骤来设置YOLOv5。


下载自定义YOLOv5对象检测数据

在本教程中,我们将从Roboflow下载YOLOv5格式的自定义对象检测数据。在本教程中,我们使用公共血细胞检测数据集训练YOLOv5检测血流中的细胞,你可以使用公共血细胞数据集或上传你自己的数据集。
  • Roboflow:https://roboflow.ai/
  • 公共血细胞数据集:https://public.roboflow.ai/object-detection/bccd


关于标记工具的说明

如果你有未标记的图像,则首先需要标记它们。对于免费的开源标签工具,我们推荐你阅读 LabelImg入门CVAT注释工具入门 的教程指南。尝试标记约50幅图像再继续本教程,因为在以后提高模型的性能的过程中,你将需要添加更多标签。
  • https://blog.roboflow.ai/getting-started-with-labelimg-for-labeling-object-detection-data/
  • https://blog.roboflow.ai/getting-started-with-cvat/
一旦你标记了数据,要将数据移动到Roboflow中,请创建一个免费帐户,然后你可以以任何格式拖动数据集:(VOC XML、COCO JSON、TensorFlow对象检测CSV等)。
上传后,你可以选择预处理和增强步骤:

为BCCD示例数据集选择设置
然后,单击 Generate 和 Download,你将可以选择YOLOv5 Pythorch格式了。

选择“YOLO v5 Pythorch”
当出现提示时,一定要选择“Show Code Snippet”,这将输出一个下载curl脚本,这样你就可以轻松地将数据以正确的格式移植到Colab中。
curl -L "https://public.roboflow.ai/ds/YOUR-LINK-HERE" > roboflow.zip; unzip roboflow.zip; rm roboflow.zip
正在Colab中下载…
下载YOLOv5格式的自定义对象数据集
导出将会创建一个名为data.yaml的YOLOv5.yaml文件,指定YOLOv5 images文件夹、YOLOv5 labels文件夹的位置以及自定义类的信息。


定义YOLOv5模型配置和架构

接下来,我们为我们的定制对象检测器编写一个模型配置文件。在本教程中,我们选择了最小、最快的YOLOv5基本模型,你也可以从其他YOLOv5模型中选择,包括:
  • YOLOv5s
  • YOLOv5m
  • YOLOv5l
  • YOLOv5x
你也可以在此步骤中编辑网络结构,但一般不需要这样做。以下是YOLOv5模型配置文件,我们将其命名为custom_yolov5s.yaml
nc: 3depth_multiple: 0.33width_multiple: 0.50
anchors: - [10,13, 16,30, 33,23] - [30,61, 62,45, 59,119] - [116,90, 156,198, 373,326]
backbone: [[-1, 1, Focus, [64, 3]], [-1, 1, Conv, [128, 3, 2]], [-1, 3, Bottleneck, [128]], [-1, 1, Conv, [256, 3, 2]], [-1, 9, BottleneckCSP, [256]], [-1, 1, Conv, [512, 3, 2]], [-1, 9, BottleneckCSP, [512]], [-1, 1, Conv, [1024, 3, 2]], [-1, 1, SPP, [1024, [5, 9, 13]]], [-1, 6, BottleneckCSP, [1024]], ]
head: [[-1, 3, BottleneckCSP, [1024, False]], [-1, 1, nn.Conv2d, [na * (nc + 5), 1, 1, 0]], [-2, 1, nn.Upsample, [None, 2, "nearest"]], [[-1, 6], 1, Concat, [1]], [-1, 1, Conv, [512, 1, 1]], [-1, 3, BottleneckCSP, [512, False]], [-1, 1, nn.Conv2d, [na * (nc + 5), 1, 1, 0]], [-2, 1, nn.Upsample, [None, 2, "nearest"]], [[-1, 4], 1, Concat, [1]], [-1, 1, Conv, [256, 1, 1]], [-1, 3, BottleneckCSP, [256, False]], [-1, 1, nn.Conv2d, [na * (nc + 5), 1, 1, 0]],
[[], 1, Detect, [nc, anchors]], ]

训练定制YOLOv5探测器

我们的data.yamlcustom_yolov5s.yaml文件已经准备好了,我们库开始训练了!
为了开始训练,我们使用以下选项运行训练命令:
  • img:定义输入图像大小
  • batch:确定batch大小
  • epochs:定义epochs。(注:通常,3000+很常见!)
  • data:设置yaml文件的路径
  • cfg:指定我们的模型配置
  • weights:指定权重的自定义路径。(注意:你可以从Ultralytics Google Drive文件夹下载权重)
  • name:结果名称
  • nosave:只保存最后的检查点
  • cache:缓存图像以加快训练速度
运行训练命令:

训练定制的YOLOv5探测器。它训练得很快!
在训练期间,你可以看 mAP@0.5 来了解你的探测器是如何运行的,请参阅这篇文章。
  • https://blog.roboflow.ai/what-is-mean-average-precision-object-detection/


评估定制YOLOv5检测器的性能

现在我们已经完成了训练,我们可以通过查看验证指标来评估训练过程的执行情况。训练脚本将删除tensorboard日志,我们将其可视化:
在我们的自定义数据集上可视化tensorboard结果
如果你因为一些原因不能把张量可视化,可以用utils.plot_result来绘制并保存为result.png。
你需要在验证评估分数达到其最高点处获取训练好的模型权重。


可视化YOLOv5训练数据

在训练过程中,我们可以可视化真实训练数据和增强后的训练数据。
我们的真实训练数据
我们的训练数据采用自动YOLOv5增强


对测试图像运行YOLOv5推断

现在我们利用我们训练好的模型,对测试图像进行推理。训练完成后,模型权重将保存到 weights/。
推理过程,我们调用这些权重和一个指定模型置信度的conf(要求的置信度越高,预测越少)、以及一个推理源。源可以接受一个包含图像、单个图像、视频文件以及设备的网络摄像头端口的目录。对于源代码,我将test/*jpg移到test-infer/
!python detect.py --weights weights/last_yolov5s_custom.pt --img 416 --conf 0.4 --source ../test_infer
推理时间非常快,在我们的 Tesla P100 上,YOLOv5s 达到了每秒142帧!!
以142 FPS(0.007s/图像)的速度推断YOLOv5s
最后,我们在测试图像上可视化我们的检测器推断结果。
测试图像的YOLOv5推理


导出并保存YOLOv5权重以供将来推断

既然我们定制的YOLOv5物体检测器已经通过验证,我们需要从Colab中取出权重,用于实时计算机视觉任务,为此我们导入一个Google驱动器模块并将其发送出去。
from google.colab import drivedrive.mount('/content/gdrive')
%cp /content/yolov5/weights/last_yolov5s_custom.pt /content/gdrive/My\ Drive


结论

我们希望你可以训练属于你自己的定制YOLOv5检测器!
使用 YOLOv5 非常方便,而且训练迅速,推理迅速,表现出色。让我们把它弄出来!
参考链接:https://blog.roboflow.ai/how-to-train-yolov5-on-a-custom-dataset/

下载1:动手学深度学习


AI算法与图像处公众号后台回复:动手学深度学习,即可下载547页《动手学深度学习》电子书和源码。该书是面向中文读者的能运行、可讨论的深度学习教科书,它将文字、公式、图像、代码和运行结果结合在一起。本书将全面介绍深度学习从模型构造到模型训练,以及它们在计算机视觉和自然语言处理中的应用。



下载2
AI算法与图像处公众号后台回复:OpenCV实战项目20讲即可下载20个有趣的OpenCV实战项目
个人微信(如果没有备注不拉群!
请注明:地区+学校/企业+研究方向+昵称

浏览 43
点赞
评论
收藏
分享

手机扫一扫分享

举报
评论
图片
表情
推荐
点赞
评论
收藏
分享

手机扫一扫分享

举报