使用YOLOv6训练自定义数据集完整指南

温欣晶Eve

于 2025-06-07 09:07:13 发布

阅读量304

点赞数 3

CC 4.0 BY-SA版权

本文链接：https://blog.youkuaiyun.com/gitblog_00373/article/details/148488864

使用YOLOv6训练自定义数据集完整指南

YOLOv6 meituan/YOLOv6: 是一个由美团点评团队开发的YOLO系列目标检测模型。适合用于需要高性能目标检测的应用。特点是可以提供优化的网络结构和训练流程，以达到更高的检测准确率和速度。项目地址: https://gitcode.com/gh_mirrors/yo/YOLOv6

前言

YOLOv6作为美团开源的优秀目标检测算法，在实际应用中经常需要针对特定场景进行定制化训练。本文将详细介绍如何使用YOLOv6训练自定义数据集的全流程，帮助开发者快速上手并应用于实际项目中。

环境准备

在开始训练前，需要确保满足以下条件：

Python 3.8环境
已安装PyTorch框架（建议1.8+版本）
安装YOLOv6所需依赖库

建议使用conda创建独立环境以避免依赖冲突。

数据集准备

1. 数据标注

训练自定义数据集的第一步是准备标注数据。推荐使用以下工具进行标注：

LabelMe：适合小规模数据标注
Roboflow：提供在线标注和数据集管理功能

2. 标注格式转换

YOLOv6要求标注文件为YOLO格式，每个图像对应一个.txt文件，格式如下：

<类别ID> <中心点x坐标> <中心点y坐标> <框宽度> <框高度>

注意：

坐标值需要归一化到0-1之间
类别ID从0开始计数
每个对象占一行

3. 目录结构组织

规范的目录结构对训练至关重要，建议按以下方式组织：

custom_dataset/
├── images/
│   ├── train/    # 训练集图片
│   ├── val/      # 验证集图片
│   └── test/     # 测试集图片(可选)
└── labels/
    ├── train/    # 训练集标签
    ├── val/      # 验证集标签
    └── test/     # 测试集标签(可选)

4. 创建数据集配置文件

在YOLOv6的data目录下创建dataset.yaml文件，配置内容示例：

train: ../custom_dataset/images/train
val: ../custom_dataset/images/val
test: ../custom_dataset/images/test

is_coco: False  # 非COCO数据集设为False

nc: 3  # 类别数量
names: ['cat', 'dog', 'person']  # 类别名称

模型配置

YOLOv6提供了灵活的配置系统，可以通过修改配置文件调整模型结构和训练参数。

1. 配置文件详解

典型配置文件包含三个主要部分：

# 模型结构配置
model = dict(
    type='YOLOv6s',  # 模型类型
    pretrained='./weights/yolov6s.pt',  # 预训练权重
    depth_multiple=0.33,  # 深度系数
    width_multiple=0.50,  # 宽度系数
    ...
)

# 优化器配置
solver=dict(
    optim='SGD',  # 优化器类型
    lr_scheduler='Cosine',  # 学习率调度策略
    ...
)

# 数据增强配置
data_aug = dict(
    hsv_h=0.015,  # 色调增强幅度
    hsv_s=0.7,    # 饱和度增强幅度
    hsv_v=0.4,    # 明度增强幅度
    ...
)

2. 预训练模型选择

根据需求选择合适的预训练模型：

YOLOv6s：轻量级，适合移动端
YOLOv6m：平衡型
YOLOv6l：高性能

训练过程

1. 单GPU训练

基础训练命令：

python tools/train.py \
    --batch 32 \
    --conf configs/yolov6s_finetune.py \
    --data data/dataset.yaml \
    --fuse_ab \
    --device 0

参数说明：

--batch：批次大小
--conf：配置文件路径
--data：数据集配置文件
--fuse_ab：启用Anchor-based和Anchor-free融合训练
--device：指定GPU设备

2. 多GPU分布式训练

python -m torch.distributed.launch \
    --nproc_per_node 8 \
    tools/train.py \
    --batch 256 \
    --conf configs/yolov6s_finetune.py \
    --data data/dataset.yaml \
    --fuse_ab \
    --device 0,1,2,3,4,5,6,7

3. 自蒸馏训练

YOLOv6支持自蒸馏训练提升模型性能：

python -m torch.distributed.launch \
    --nproc_per_node 8 \
    tools/train.py \
    --batch 256 \
    --conf configs/yolov6s_finetune.py \
    --data data/dataset.yaml \
    --distill \
    --teacher_model_path your_model_path \
    --device 0,1,2,3,4,5,6,7

注意：自蒸馏训练需要先在配置文件中启用DFL（Distribution Focal Loss）：

use_dfl=True
reg_max=16

模型评估

训练完成后，使用以下命令评估模型性能：

python tools/eval.py \
    --data data/data.yaml \
    --weights output_dir/name/weights/best_ckpt.pt \
    --task val \
    --device 0

模型推理

使用训练好的模型进行预测：

python tools/infer.py \
    --weights output_dir/name/weights/best_ckpt.pt \
    --source img.jpg \
    --device 0

支持多种输入源：

单张图片
图片目录
视频文件
摄像头输入

模型部署

YOLOv6支持多种部署方式，以下介绍ONNX导出：

1. 基础ONNX导出

python deploy/ONNX/export_onnx.py \
    --weights output_dir/name/weights/best_ckpt.pt \
    --simplify \
    --device 0

2. 带NMS的ONNX导出（ONNX Runtime）

python deploy/ONNX/export_onnx.py \
    --weights output_dir/name/weights/best_ckpt.pt \
    --simplify \
    --device 0 \
    --dynamic-batch \
    --end2end \
    --ort

3. 带NMS的ONNX导出（TensorRT）

python deploy/ONNX/export_onnx.py \
    --weights output_dir/name/weights/best_ckpt.pt \
    --simplify \
    --device 0 \
    --dynamic-batch \
    --end2end