使用 YOLOv5 训练自定义目标检测模型

一、环境准备

1. 克隆 YOLOv5 仓库

首先，我们需要从 GitHub 上克隆 YOLOv5 仓库到本地：

git clone https://github.com/ultralytics/yolov5
cd yolov5

2. 安装依赖

进入克隆下来的 yolov5 目录后，使用以下命令安装所需的依赖库：

pip install -r requirements.txt

二、数据集准备

1. 数据集格式

YOLOv5 要求数据集按照特定的格式进行组织，数据集应包含图像文件和对应的标注文件。标注文件采用 YOLO 格式，每一行代表一个目标，格式为 class_id x_center y_center width height，其中坐标和宽高都是相对于图像宽度和高度的比例值。

2. 数据集目录结构

数据集应按照以下结构进行组织：

dataset/
├── images/
│   ├── train/
│   │   ├── image_001.jpg
│   │   ├── image_002.jpg
│   │   └── ...
│   ├── val/
│   │   ├── image_051.jpg
│   │   ├── image_052.jpg
│   │   └── ...
│   └── test/
│       ├── image_101.jpg
│       ├── image_102.jpg
│       └── ...
├── labels/
│   ├── train/
│   │   ├── image_001.txt
│   │   ├── image_002.txt
│   │   └── ...
│   ├── val/
│   │   ├── image_051.txt
│   │   ├── image_052.txt
│   │   └── ...
│   └── test/
│       ├── image_101.txt
│       ├── image_102.txt
│       └── ...
├── data.yaml
├── README.md
└── logs/
    ├── data_preprocessing.log
    └── annotation_check.log

• images 文件夹用于存放所有的图像文件，根据用途分为 train（训练集）、val（验证集）和 test（测试集）三个子文件夹。

• labels 文件夹用于存放对应的标注文件，同样分为 train、val 和 test 三个子文件夹，标注文件的文件名与对应的图像文件名相同，扩展名为 .txt。

3. 创建数据集配置文件

在 data 文件夹下创建一个 .yaml 文件，例如 custom_data.yaml，内容如下：

train: ../dataset/images/train/
val: ../dataset/images/val/
test: ../dataset/images/test/

nc: 2  # 类别数量
names: ['class1', 'class2']  # 类别名称

三、选择模型配置文件

YOLOv5 提供了不同大小的预训练模型，如 yolov5s.yaml、yolov5m.yaml、yolov5l.yaml 和 yolov5x.yaml，你可以根据自己的需求选择合适的模型。模型配置文件位于 models 文件夹下。

四、训练模型

使用 train.py 脚本进行模型训练，以下是一个示例命令：

python train.py --img 640 --batch 16 --epochs 100 --data data/custom_data.yaml --cfg models/yolov5s.yaml --weights yolov5s.pt

参数说明：

• --img：输入图像的大小，这里设置为 640x640。

• --batch：批量大小，根据你的 GPU 内存进行调整。

• --epochs：训练的轮数。

• --data：数据集配置文件的路径。

• --cfg：模型配置文件的路径。

• --weights：预训练模型的权重文件路径。

五、训练过程监控

在训练过程中，YOLOv5 会在 runs/train 文件夹下创建一个新的文件夹，用于保存训练日志、模型权重和训练结果。你可以使用 TensorBoard 来监控训练过程：

tensorboard --logdir runs/train

六、模型评估

训练完成后，你可以使用 val.py 脚本来评估模型在测试集上的性能：

python val.py --weights runs/train/exp/weights/best.pt --data data/custom_data.yaml --img 640

参数说明：

• --weights：训练好的模型权重文件路径。

• --data：数据集配置文件的路径。

• --img：输入图像的大小。

七、模型推理

使用训练好的模型进行推理，你可以使用 detect.py 脚本：

python detect.py --weights runs/train/exp/weights/best.pt --source path/to/images/or/video

参数说明：

• --weights：训练好的模型权重文件路径。

• --source：输入图像或视频的路径。

八、常见问题及解决方法

1. GPU 内存不足

如果在训练过程中出现 GPU 内存不足的问题，可以尝试减小 --batch 参数的值。

2. 数据集路径错误

确保 custom_data.yaml 文件中的数据集路径正确，并且数据集目录结构符合要求。

3. 训练速度慢

可以尝试使用更强大的 GPU 或者调整 --batch 和 --img 参数的值。

九、总结

通过以上步骤，你就可以使用 YOLOv5 训练自己的目标检测模型了。希望本指南对你有所帮助！