📚 课程目录第2章自定义数据集训练YOLOv8检测模型

2. 自定义数据集训练YOLOv8检测模型

任务目标

本案例将使用自定义标注数据集（Pascal VOC格式）训练YOLOv8检测模型，完成模型评估，并导出可用于后续DeepSORT多目标跟踪的检测权重。通过本案例，你将掌握从自定义数据准备到模型部署的完整流程。

环境与依赖

Python 3.8+
ultralytics（YOLOv8官方库）：pip install ultralytics
PyTorch 1.8+（根据CUDA版本安装）
数据集标注工具：LabelImg或LabelStudio（用于生成VOC格式标注）

数据集准备

本案例以车辆检测为例，数据集目录结构如下（需自行准备或使用公开车辆数据集）：

datasets/
└── vehicle_data/
    ├── images/
    │   ├── train/    # 训练图片（.jpg）
    │   └── val/      # 验证图片（.jpg）
    ├── labels/
    │   ├── train/    # 训练标注（.xml，VOC格式）
    │   └── val/      # 验证标注（.xml）
    └── data.yaml     # 数据集配置文件

首先将VOC格式标注转换为YOLO格式，并生成data.yaml配置文件：

import os
import xml.etree.ElementTree as ET
from pathlib import Path

def voc_to_yolo(xml_path, out_dir, class_names):
    """将单个VOC XML标注转换为YOLO txt格式"""
    tree = ET.parse(xml_path)
    root = tree.getroot()
    img_w = int(root.find('size/width').text)
    img_h = int(root.find('size/height').text)
    
    txt_name = Path(xml_path).stem + '.txt'
    with open(os.path.join(out_dir, txt_name), 'w') as f:
        for obj in root.findall('object'):
            cls = obj.find('name').text
            if cls not in class_names:
                continue
            cls_id = class_names.index(cls)
            bbox = obj.find('bndbox')
            x1 = float(bbox.find('xmin').text)
            y1 = float(bbox.find('ymin').text)
            x2 = float(bbox.find('xmax').text)
            y2 = float(bbox.find('ymax').text)
            
            # 归一化到[0,1]
            x_center = (x1 + x2) / 2 / img_w
            y_center = (y1 + y2) / 2 / img_h
            w = (x2 - x1) / img_w
            h = (y2 - y1) / img_h
            f.write(f"{cls_id} {x_center:.6f} {y_center:.6f} {w:.6f} {h:.6f}\n")

# 类别定义（按需修改）
CLASS_NAMES = ['car', 'truck', 'bus', 'motorcycle']

# 转换训练集和验证集
for split in ['train', 'val']:
    xml_dir = f'datasets/vehicle_data/labels/{split}'
    out_dir = f'datasets/vehicle_data/labels_yolo/{split}'
    os.makedirs(out_dir, exist_ok=True)
    for xml_file in os.listdir(xml_dir):
        if xml_file.endswith('.xml'):
            voc_to_yolo(os.path.join(xml_dir, xml_file), out_dir, CLASS_NAMES)

创建data.yaml配置文件：

# datasets/vehicle_data/data.yaml
path: datasets/vehicle_data  # 数据集根目录
train: images/train          # 训练图片相对路径
val: images/val              # 验证图片相对路径

nc: 4  # 类别数量
names: ['car', 'truck', 'bus', 'motorcycle']  # 类别名称

模型训练与评估

以下代码完成YOLOv8模型的训练、验证和导出全流程：

from ultralytics import YOLO
import torch

def main():
    # 1. 加载预训练模型（推荐使用COCO预训练权重进行迁移学习）
    model = YOLO('yolov8n.pt')  # 可选：yolov8s.pt, yolov8m.pt等
    
    # 2. 训练模型
    results = model.train(
        data='datasets/vehicle_data/data.yaml',
        epochs=100,           # 训练轮数
        imgsz=640,            # 输入图片尺寸
        batch=16,             # 批次大小（根据GPU显存调整）
        device=0,             # GPU设备ID，CPU则设为'cpu'
        workers=4,            # 数据加载线程数
        lr0=0.01,             # 初始学习率
        patience=20,          # 早停耐心值
        project='runs/train', # 训练结果保存路径
        name='vehicle_det',   # 实验名称
        exist_ok=True         # 允许覆盖同名实验
    )
    
    # 3. 在验证集上评估模型性能
    metrics = model.val(
        data='datasets/vehicle_data/data.yaml',
        split='val',
        project='runs/val',
        name='vehicle_det_eval',
        exist_ok=True
    )
    
    # 打印关键评估指标
    print(f"mAP50: {metrics.box.map50:.4f}")
    print(f"mAP50-95: {metrics.box.map:.4f}")
    print(f"Precision: {metrics.box.mp:.4f}")
    print(f"Recall: {metrics.box.mr:.4f}")
    
    # 4. 导出模型为TorchScript格式（用于DeepSORT跟踪）
    best_model = YOLO('runs/train/vehicle_det/weights/best.pt')
    exported_path = best_model.export(format='torchscript', imgsz=640)
    print(f"模型已导出至: {exported_path}")
    
    # 同时导出ONNX格式（可选，便于跨平台部署）
    onnx_path = best_model.export(format='onnx', imgsz=640)
    print(f"ONNX模型已导出至: {onnx_path}")

if __name__ == '__main__':
    main()

运行方式

将上述代码保存为train_yolov8.py，确保数据集目录结构正确；
执行训练命令：python train_yolov8.py；
训练过程中可通过TensorBoard监控：tensorboard --logdir runs/train。

验证方法与通过标准

训练收敛验证：训练完成后，检查runs/train/vehicle_det/目录下的results.png，确认loss曲线下降且无明显过拟合（训练/验证loss差距不大）；
评估指标验证：验证集上mAP50应达到0.85以上（若数据集质量良好且类别数≤10），若低于0.7需检查标注质量或增加训练轮数；
可视化验证：运行以下命令随机抽取验证集图片查看检测效果：

from ultralytics import YOLO
model = YOLO('runs/train/vehicle_det/weights/best.pt')
results = model.predict(source='datasets/vehicle_data/images/val', save=True, project='runs/predict', name='val_samples')

导出验证：确认best.torchscript文件生成，且能被后续DeepSORT案例正常加载（加载测试：torch.jit.load('runs/train/vehicle_det/weights/best.torchscript')无报错）。

常见问题

CUDA内存不足：降低batch值或imgsz，或使用梯度累积；
类别不匹配：确保data.yaml中nc和names与标注转换脚本中的CLASS_NAMES完全一致；
标注文件缺失：YOLO格式要求每张图片对应一个同名txt文件，若图片无目标则生成空txt文件。