📚 课程目录第2章 YOLOv5训练自定义数据集

2. YOLOv5训练自定义数据集

任务目标

本案例将使用一份小型自定义口罩佩戴检测数据集，完成YOLOv5模型的端到端训练。你将学会：准备YOLO格式数据集、配置模型与超参数、启动训练并监控收敛过程，最终获得可用于推理的权重文件。

环境与依赖

需要一台具备NVIDIA GPU（建议显存≥8GB）的Linux/Windows机器，并安装以下依赖：

# Python 3.8+
pip install torch==1.13.1 torchvision==0.14.1 --extra-index-url https://download.pytorch.org/whl/cu117
pip install opencv-python pyyaml tqdm matplotlib seaborn
git clone https://github.com/ultralytics/yolov5.git
cd yolov5
pip install -r requirements.txt

数据集准备

使用公开的Mask Wearing Dataset（包含853张图片，两类：with_mask / without_mask）。将数据集按8:2划分为训练集与验证集，并转换为YOLO格式（每张图片对应一个同名.txt文件，每行：class_id x_center y_center width height，坐标为归一化后的值）。

数据集目录结构如下：

datasets/mask/
├── images/
│   ├── train/   (682张)
│   └── val/     (171张)
├── labels/
│   ├── train/   (对应txt)
│   └── val/
└── mask.yaml

mask.yaml内容：

train: datasets/mask/images/train
val: datasets/mask/images/val
nc: 2
names: ['with_mask', 'without_mask']

关键完整代码

以下脚本完成训练前的数据校验与训练启动。将以下内容保存为train_mask.py，放在yolov5目录下运行：

import torch
from pathlib import Path
import yaml

# 1. 校验数据集完整性
data_cfg = "datasets/mask/mask.yaml"
with open(data_cfg) as f:
    data = yaml.safe_load(f)

for split in ["train", "val"]:
    img_dir = Path(data[split])
    lbl_dir = img_dir.parent.parent / "labels" / split
    assert img_dir.exists(), f"图片目录缺失: {img_dir}"
    assert lbl_dir.exists(), f"标签目录缺失: {lbl_dir}"
    imgs = list(img_dir.glob("*.jpg")) + list(img_dir.glob("*.png"))
    assert len(imgs) > 0, f"{split}集为空"
    # 抽查前5张是否有对应标签
    for img in imgs[:5]:
        lbl = lbl_dir / (img.stem + ".txt")
        assert lbl.exists(), f"缺少标签: {lbl}"

print("数据集校验通过，开始训练...")

# 2. 训练超参数（根据GPU显存调整）
hyp = {
    "lr0": 0.01,          # 初始学习率
    "batch_size": 16,     # 显存不足可降至8
    "epochs": 100,
    "img_size": 640,
    "workers": 4,
    "device": 0,          # GPU编号
}

# 3. 启动训练（使用yolov5s预训练权重）
torch.cuda.empty_cache()
import train as yolov5_train

yolov5_train.main(
    hyp="data/hyps/hyp.scratch-low.yaml",  # 使用官方基础超参文件
    data=data_cfg,
    weights="yolov5s.pt",   # 自动下载COCO预训练权重
    batch_size=hyp["batch_size"],
    epochs=hyp["epochs"],
    imgsz=hyp["img_size"],
    workers=hyp["workers"],
    device=hyp["device"],
    project="runs/train",
    name="mask_exp",
    exist_ok=True,
)

运行与监控

首次运行会自动下载yolov5s.pt预训练权重（约14MB）。
训练日志实时打印：Epoch、GPU_mem、box_loss、cls_loss、mAP@0.5等指标。
每轮结束在runs/train/mask_exp/下生成results.png，包含loss曲线与mAP曲线。

验证方法与通过标准

训练完成后（约30-60分钟，取决于GPU），执行以下验证：

python detect.py --weights runs/train/mask_exp/weights/best.pt \
                 --source datasets/mask/images/val/ --conf 0.5

通过条件：

训练日志中mAP@0.5最终达到0.90以上（该数据集简单，通常可达0.95+）。
detect.py输出的预测图片中，口罩佩戴目标被正确框出，且类别标签正确。
results.png中loss曲线呈下降趋势且无明显发散。

常见问题

CUDA out of memory：将batch_size降至8或4，或减小img_size至512。
标签格式错误：确保txt中坐标为归一化浮点数（0~1），且类别ID在nc范围内。
mAP不涨：检查学习率是否过大（lr0可降至0.005），或训练轮次不足（增加至150轮）。

训练完成后，best.pt即为可用于后续推理或迁移学习的最优权重。