从零到一：使用YOLOv5训练自定义数据集

引言

YOLOv5（You Only Look Once version 5）是当前最流行且高效的目标检测模型之一，以其快速、准确的特点受到广大开发者的青睐。本文将指导你如何从头开始，使用YOLOv5训练自己的数据集。

1. 环境搭建

首先，确保你的开发环境已经安装了Python和必要的库。YOLOv5通常依赖PyTorch框架，以及其他一些Python库如numpy、opencv等。

步骤一：安装PyTorch

根据你的操作系统和CUDA版本，访问PyTorch官网获取安装命令。例如，如果你使用CUDA 11.3，可以运行类似以下命令：

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

步骤二：克隆YOLOv5仓库

使用git克隆YOLOv5的GitHub仓库，或者直接下载ZIP包解压。

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

步骤三：安装YOLOv5依赖

在YOLOv5目录下，运行以下命令安装依赖库。

pip install -r requirements.txt

2. 数据集准备

步骤一：收集图片

收集你需要进行目标检测的图片，并整理到一个文件夹中。建议图片格式统一为JPG。

步骤二：标注图片

使用标注工具（如LabelImg）对图片中的目标进行标注。标注时，需要为每个目标绘制一个检测框，并指定其类别。

YOLOv5要求标注数据为txt格式，包含类别ID、归一化的中心点坐标（x, y）以及宽高（w, h）。标注完成后，将生成的txt文件与对应的图片放在同一目录下。

步骤三：划分数据集

将标注好的数据集划分为训练集、验证集和测试集。可以使用Python脚本来完成这一任务，确保每个子集的数据分布均匀。

3. 配置YOLOv5

步骤一：修改配置文件

在YOLOv5的data目录下，创建一个新的YAML文件（如mydata.yaml），用于指定数据集的路径、类别等信息。

# mydata.yaml
train: data/mydata/images/train  # 训练集图片路径
val: data/mydata/images/val      # 验证集图片路径
# number of classes
nc: 2  # 假设你有两个类别
names: ['class1', 'class2']  # 类别名称

步骤二：选择预训练模型

在YOLOv5的models目录下，选择一个合适的预训练模型（如yolov5s.pt），并准备好其路径。

4. 训练模型

使用YOLOv5提供的train.py脚本来训练模型。通过命令行传递必要的参数，如预训练模型路径、配置文件路径等。

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

这里的--img 640指定了输入图片的大小，--batch 16是批处理大小，--epochs 50是训练轮次，--data指定了数据集配置文件，--weights指定了预训练模型路径。

5. 模型评估与测试

训练完成后，可以使用YOLOv5提供的val.py脚本来评估模型在验证集上的性能。

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

此外，你还可以使用detect.py脚本来对新的图片或视频进行目标检测，以验证模型的泛化能力。