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

点击上方“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预装了 torch 和 cuda 。若是你尝试在本地使用本教程，可能须要执行其余步骤来设置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.yaml 和 custom_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实战项目

    
 
 
     我的微信（若是没有备注不拉群！） 

    

    
    
    
     
     
              
     

    
 
 
     请注明： 
 
     地区+学校/企业+研究方向+昵称 

    

    
    
    
     
     
              
     

    
 
 
     
 

    

本文分享自微信公众号 - AI算法与图像处理（AI_study）。
若有侵权，请联系 support@oschina.cn 删除。
本文参与“OSC源创计划”，欢迎正在阅读的你也加入，一块儿分享。