使用YOLOv5训练自定义数据集

下面一步步说明如何使用 YOLOv5 训练它，得到一个可以进行预测的目标检测模型。

1. 准备工作

在训练之前，需要确保计算机上安装了必要的工具和库。以下是需要准备的：

• Python 3.7 或更高版本：YOLOv5 依赖 Python。
• PyTorch：YOLOv5 是基于 PyTorch 开发的。
• Git：用于下载 YOLOv5 的代码。

安装 Python

如果计算机上还没有安装 Python，可以从 Python 官网 下载并安装。安装后，在终端输入以下命令检查版本：

python --version

安装 PyTorch

访问 PyTorch 官网，根据系统（Windows/Mac/Linux）和硬件（CPU 或 GPU）选择合适的安装命令。没有NVIDIA GPU的情况下请参考使用YOLOv5训练自定义数据集-无NVIDIA GPU。如果计算机上装有 NVIDIA GPU，可以运行：

pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121

下载 YOLOv5

打开终端，运行以下命令克隆 YOLOv5 的 GitHub 仓库：

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

安装依赖项

在 yolov5 目录下，运行以下命令安装其他依赖库：

pip install -r requirements.txt

完成这些步骤后，计算机就准备好运行 YOLOv5 所需条件了！

2. 准备数据集

YOLOv5 需要图像和对应的标注文件来训练模型。自定义数据集可能已经有一些图像，但还需要确保它符合 YOLOv5 的要求。

检查自定义数据集

自定义数据集应该包含：

• 图像文件：如 .jpg、.png 等格式。
• 标注文件：每个图像需要一个同名的 .txt 文件，采用 YOLO 格式。每行格式为：

  class x_center y_center width height
  

  • class：类别编号（从 0 开始）。
  • x_center, y_center：边界框中心坐标（归一化到 0-1）。
  • width, height：边界框宽高（归一化到 0-1）。

如果没有标注文件

如果自定义的数据集还没有标注文件，可以使用工具生成。推荐 LabelImg：

1. 安装 LabelImg：

   pip install labelimg
   

2. 启动 LabelImg：

   labelimg
   

3. 打开图像文件夹，选择 “YOLO” 格式，标注对象并保存。完成后，每个图像会有一个对应的 .txt 文件。

组织数据集结构

将自定义数据集整理成以下结构：

dataset/
├── images/
│   ├── train/
│   │   ├── image1.jpg
│   │   ├── image2.jpg
│   │   └── ...
│   ├── val/
│   │   ├── image3.jpg
│   │   ├── image4.jpg
│   │   └── ...
│   └── test/
│       ├── image5.jpg
│       ├── image6.jpg
│       └── ...
└── labels/
    ├── train/
    │   ├── image1.txt
    │   ├── image2.txt
    │   └── ...
    ├── val/
    │   ├── image3.txt
    │   ├── image4.txt
    │   └── ...
    └── test/
        ├── image5.txt
        ├── image6.txt
        └── ...

• images/：存放图像，分为训练集（train）、验证集（val）和测试集（test）。
• labels/：存放标注文件，与图像同名。

数据集划分

如果自定义的数据集还没有划分数据集，可以按以下比例分配：

• 训练集：80%
• 验证集：10%
• 测试集：10%
  用户可以用手动移动文件，或者写一个简单脚本随机划分。

3. 配置数据集

YOLOv5 需要一个 .yaml 文件来指定数据集路径和类别信息。创建一个名为 data.yaml 的文件，内容如下：

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

nc: 2  # 替换为自定义数据集的类别数量（实际所需）
names: ['class1', 'class2']  # 替换为用户自定义数据集中的类别名称

• train/val/test：指向图像文件夹的路径（相对或绝对路径）。
• nc：自定义数据集中的类别数量。
• names：类别名称列表（与标注文件中的 class 对应）。

将 data.yaml 保存到 yolov5/data/ 目录下。

4. 选择模型架构

YOLOv5 提供了多种模型，从轻量级的 yolov5s 到强大的 yolov5x。对于初学者，推荐使用 yolov5s，因为它训练快且资源需求低。可以在 yolov5/models/ 目录下看到这些模型的配置文件。

5. 训练模型

现在可以开始训练了！在终端中进入 yolov5 目录，运行以下命令：

python train.py --img 640 --batch 16 --epochs 50 --data data/data.yaml --weights yolov5s.pt

参数说明

（注意：根据是否有GPU，训练参数设置有所不同，参考使用YOLOv5训练自定义数据集-无NVIDIA GPU）

• --img 640：输入图像大小（通常为 640x640）。
• --batch 16：批次大小（根据用户使用的 GPU 内存调整，例如 8 或 4）。
• --epochs 50：训练轮数（可根据数据集大小调整）。
• --data data/data.yaml：自定义配置文件的路径。
• --weights yolov5s.pt：预训练权重（从 yolov5s 开始）。

训练过程会显示损失值等信息，模型权重会自动保存到 runs/train/exp/ 目录。

常见问题解决

• 内存不足：减少 --batch 大小或使用更小的模型。
• 路径错误：检查 data.yaml 中的路径是否正确。
• 标注错误：确保 .txt 文件格式正确。

6. 监控训练过程

无GPU的情况下，参考使用YOLOv5训练自定义数据集-无NVIDIA GPU，防止CPU过载。训练时，YOLOv5 会生成日志。可以用 TensorBoard 查看损失曲线和指标：

tensorboard --logdir runs/train

在浏览器打开 http://localhost:6006 查看。

7. 评估模型

训练完成后，评估模型性能。运行：

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

这会输出 mAP（平均精度均值），越高越好。

8. 使用模型预测

现在可以用模型预测新图像了！运行：

python detect.py --weights runs/train/exp/weights/best.pt --img 640 --conf 0.25 --source path/to/your/image.jpg

• --conf 0.25：置信度阈值。
• --source：图像或视频路径。

结果会保存在 runs/detect/exp/ 目录下。

9. 优化模型（可选）

如果结果不理想，可以尝试：

• 增加数据量。
• 调整 --epochs 或 --batch。
• 使用更大的模型（如 yolov5m）。

总结

以上就是用 YOLOv5 训练了用户自定义的数据集，并得到了一个预测模型的步骤。简单回顾如下：

1. 安装环境。
2. 准备和组织数据集。
3. 配置 data.yaml。
4. 训练模型。
5. 评估和预测。