[炼丹术]YOLOv5训练自定义数据集

时间:2023-04-09 14:20:14

YOLOv5训练自定义数据

一、开始之前的准备工作

克隆 repo 并在Python>=3.6.0环境中安装requirements.txt,包括PyTorch>=1.7模型数据集会从最新的 YOLOv5版本中自动下载。

git clone https://github.com/ultralytics/yolov5
cd yolov5
pip install -r requirements.txt

二、训练自定义数据

2.1 创建my_dataset.yaml

COCO128是一个示例小教程数据集,由COCO train2017中的前 128 张图像组成。这些相同的 128 张图像用于训练和验证,以验证我们的训练管道是否能够过拟合。数据/ coco128.yaml,如下所示,是数据集的配置文件,它定义1)数据集根目录path和相对路径train/ val/test图像目录(或* .txt与图像文件的路径),2)的类的数量nc和3)类列表names

# Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..]
path: ../datasets/coco128 # dataset root dir 数据集根目录
train: images/train2017 # train images (relative to 'path') 128 images #训练图像(相对于“path”)
val: images/train2017 # val images (relative to 'path') 128 images # val 图像(相对于“path”)
test: # test images (optional) #测试图像(可选) # Classes
nc: 80 # number of classes
names: [ 'person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 'truck', 'boat', 'traffic light',
'fire hydrant', 'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse', 'sheep', 'cow',
'elephant', 'bear', 'zebra', 'giraffe', 'backpack', 'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee',
'skis', 'snowboard', 'sports ball', 'kite', 'baseball bat', 'baseball glove', 'skateboard', 'surfboard',
'tennis racket', 'bottle', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple',
'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake', 'chair', 'couch',
'potted plant', 'bed', 'dining table', 'toilet', 'tv', 'laptop', 'mouse', 'remote', 'keyboard', 'cell phone',
'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 'book', 'clock', 'vase', 'scissors', 'teddy bear',
'hair drier', 'toothbrush' ] # class names

这里设置数据集路径有三种方式:

    1. dir: path/to/imgs,
    1. file: path/to/imgs.txt, 或
    1. list: [path/to/imgs1, path/to/imgs2, .. ]

2.2 创建label标签

使用CVATmakeense.ai等工具标记图像后,将标签导出为YOLO 格式*.txt每个图像一个文件(如果图像中没有对象,则不需要*.txt文件)。该*.txt文件规格有:

  • 每个对象一行
  • 每一行都是class x_center y_center width height格式。
  • 框坐标必须采用标准化 xywh格式(从 0 - 1)。如果您的箱子以像素为单位,划分x_centerwidth通过图像宽度,y_centerheight通过图像高度。
  • 类号是零索引的(从 0 开始)。

[炼丹术]YOLOv5训练自定义数据集

上图对应的标签文件包含2个人(class 0)和一条领带(class 27):

[炼丹术]YOLOv5训练自定义数据集

2.3 整理目录

根据以下示例组织您的训练和验证图像和标签。YOLOv5 假设 /coco128/datasets目录旁边/yolov5目录中。YOLOv5通过将/images/每个图像路径中的最后一个实例替换为/labels/. 例如:

../datasets/coco128/images/im0.jpg   #图像
../datasets/coco128/labels/im0.txt #标签

文件结构,如下图所示:

[炼丹术]YOLOv5训练自定义数据集

2.4 选择模型型号

选择一个预训练模型开始训练。这里我们选择YOLOv5s,这是可用的最小和最快的模型。有关所有模型的完整比较,请参阅 README

[炼丹术]YOLOv5训练自定义数据集

[炼丹术]YOLOv5训练自定义数据集

2.5 train训练

通过指定数据集、批量大小、图像大小以及预训练--weights yolov5s.pt(推荐)或随机初始化--weights '' --cfg yolov5s.yaml(不推荐),在 COCO128 上训练 YOLOv5s 模型。预训练权重是从最新的 YOLOv5 版本自动下载的。

# 在COCO128 上训练
YOLOv5s 3 epochs $ python train.py --img 640 --batch 16 --epochs 3 --data coco128.yaml --weights yolov5s.pt

所有训练结果都保存在runs/train/递增的运行目录中,即runs/train/exp2runs/train/exp3等。有关更多详细信息,请参阅我们的 Google Colab Notebook 的训练部分。

[炼丹术]YOLOv5训练自定义数据集

[炼丹术]YOLOv5训练自定义数据集

三、可视化

权重和偏差记录( 新)

权重和偏差(W&B) 现在与 YOLOv5 集成,用于训练运行的实时可视化和云记录。这允许更好地运行比较和内省,以及提高团队成员之间的可见性和协作。要启用 W&B 日志记录,请安装wandb,然后正常训练(首次使用时将指导您进行设置)。

pip install wandb

在训练过程期间,你将在https://wandb.ai看到实时更新,并且您可以使用 W&B 报告工具创建结果的详细报告

[炼丹术]YOLOv5训练自定义数据集

四、本地日志

所有的结果都在默认情况下记录runs/train,为每个新的培训作为创建一个新的实验目录runs/train/exp2runs/train/exp3等查看火车和Val JPG文件看马赛克,标签,预测和增强效果。请注意,使用 Ultralytics Mosaic Dataloader进行训练(如下所示),它在训练期间将 4 个图像组合成 1 个马赛克。

train_batch0.jpg 显示训练批次 0 马赛克和标签:

[炼丹术]YOLOv5训练自定义数据集

val_batch0_labels.jpg 显示 val 批次 0 标签:

[炼丹术]YOLOv5训练自定义数据集

val_batch0_pred.jpg显示 val 批次 0预测

[炼丹术]YOLOv5训练自定义数据集

训练结果自动记录到TensorboardCSVresults.csvresults.png训练完成后绘制为(下图)。您还可以results.csv手动绘制任何文件:

from utils.plots import plot_results
plot_results('path/to/results.csv') # plot 'results.csv' as 'results.png'

[炼丹术]YOLOv5训练自定义数据集