YOLOv5 实战:从 GitHub 拉取到自定义数据集训练

YOLOv5 实战:从 GitHub 拉取到自定义数据集训练

系列说明:本篇为 YOLO 系列第四篇,聚焦 YOLOv5 的工程化实践------从 GitHub 拉取代码、准备 MaskDataSet 数据集,到完成模型训练与推理的完整流程。


一、YOLOv5 简介

GitHub 仓库https://github.com/ultralytics/yolov5
作者/组织 :Ultralytics
最新版本:v7.0(支持目标检测 + 实例分割 + 图像分类)

YOLOv5 与 YOLOv4 的核心区别在于工程化程度

对比项 YOLOv4 YOLOv5
实现框架 Darknet(纯C) PyTorch(主流深度学习框架)
代码质量 学术风格 工程化完善(模块化、自动化)
部署支持 需手动转换 原生支持 ONNX/TensorRT/CoreML/TFLite
训练体验 配置复杂 参数简洁、功能丰富
社区生态 一般 极其活跃(Star 4万+)

💡 建议:学习目标检测首选 YOLOv5,生产部署可参考 YOLOv5 工程结构。


二、环境准备

2.1 系统要求

组件 最低要求 推荐配置
Python ≥ 3.8.0 3.9 / 3.10
PyTorch ≥ 1.8 2.0+
GPU 可选(CPU可训练但很慢) NVIDIA GPU + CUDA 11.x
内存 8GB 16GB+
显存 - 8GB+(训练 YOLOv5s)

2.2 安装 PyTorch(GPU 版本)

bash 复制代码
# CUDA 11.8
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118

# CUDA 11.8 + cuDNN
pip install torch==2.0.0 torchvision==0.15.0 --index-url https://download.pytorch.org/whl/cu118

2.3 验证 PyTorch 安装

python 复制代码
import torch
print(f"PyTorch 版本: {torch.__version__}")
print(f"CUDA 可用: {torch.cuda.is_available()}")
print(f"CUDA 版本: {torch.version.cuda}")
# 输出示例:
# PyTorch 版本: 2.0.0
# CUDA 可用: True
# CUDA 版本: 11.8

三、从 GitHub 拉取 YOLOv5

3.1 克隆仓库

bash 复制代码
# 方法一:Git 克隆(推荐)
git clone https://github.com/ultralytics/yolov5.git
cd yolov5

# 方法二:指定版本
git clone --branch v7.0 https://github.com/ultralytics/yolov5.git
cd yolov5

3.2 安装依赖

bash 复制代码
# 进入项目目录
cd yolov5

# 安装所有依赖
pip install -r requirements.txt

requirements.txt 主要依赖:

txt 复制代码
以YOLOv5 的v7.0版本为例

3.3 项目结构

复制代码
yolov5/
├── train.py              # 训练脚本(核心)
├── val.py                # 验证脚本
├── detect.py             # 推理/检测脚本
├── export.py             # 模型导出脚本
├── segment/              # 实例分割模块
│   ├── train.py
│   ├── val.py
│   └── predict.py
├── classify/              # 图像分类模块
│   ├── train.py
│   ├── val.py
│   └── predict.py
├── models/               # 模型架构定义
│   ├── yolov5n.yaml
│   ├── yolov5s.yaml
│   ├── yolov5m.yaml
│   ├── yolov5l.yaml
│   └── yolov5x.yaml
├── data/                 # 数据集配置
│   ├── images/            # 数据集图片(可选)
│   ├── coco.yaml         # COCO 数据集配置
│   ├── coco128.yaml      # COCO128 小数据集配置
│   └── hyps/             # 超参数文件
├── utils/                # 工具函数
├── runs/                 # 训练输出目录
│   └── train/
│       └── exp/          # 每次训练的结果
└── tutorial.ipynb        # Jupyter Notebook 教程

四、数据集准备(MaskDataSet)

4.1 数据集格式要求

YOLOv5 使用 YOLO 格式的标注文件,目录结构如下:

复制代码
MaskDataSet/                    # 数据集根目录
├── images/
│   ├── train/                  # 训练图片
│   │   ├── img001.jpg
│   │   ├── img002.jpg
│   │   └── ...
│   ├── val/                    # 验证图片
│   │   ├── img101.jpg
│   │   └── ...
│   └── test/                   # 测试图片(可选)
│       └── ...
└── labels/
    ├── train/                  # 训练标签(与图片一一对应)
    │   ├── img001.txt
    │   ├── img002.txt
    │   └── ...
    ├── val/
    │   ├── img101.txt
    │   └── ...
    └── test/
        └── ...

4.2 标签格式说明

每个 .txt 文件对应一张图片,每行描述一个目标:

复制代码
<class_id> <x_center> <y_center> <width> <height>
  • <class_id>:类别索引(从 0 开始)
  • <x_center><y_center>:目标中心点相对图像的归一化坐标(0~1)
  • <width><height>:目标宽高相对图像的归一化值(0~1)

示例 (假设 with_mask=0,without_mask=1,mask_weared_incorrect=2):

复制代码
# img001.jpg 对应的 img001.txt
0 0.4167 0.3125 0.1667 0.1875   # 第一个目标:with_mask
1 0.6875 0.4375 0.1250 0.1563   # 第二个目标:without_mask

⚠️ 注意 :如果图像中没有目标,.txt 文件为空或不存在该文件。

4.3 使用 LabelImg 标注数据集

如果还没有标注数据,可以使用 LabelImg 工具进行标注:

bash 复制代码
# 安装 LabelImg
pip install labelimg

# 启动(设置为 YOLO 格式)
labelimg --BoundingBoxDir dataset/images/train --SaveDir dataset/labels/train --Format YOLO

标注步骤:

  1. 打开图片文件夹
  2. W 键创建矩形框
  3. 绘制目标区域,输入类别名称
  4. D 键下一张,A 键上一张
  5. 保存后自动生成同名的 .txt 文件

图1:使用 LabelImg 进行目标检测标注

4.4 数据集划分脚本

python 复制代码
import os
import shutil
import random
from pathlib import Path

def split_dataset(image_dir, output_dir, train_ratio=0.8, val_ratio=0.1):
    """
    将数据集划分为训练集、验证集、测试集
    
    Args:
        image_dir: 原始图片目录
        output_dir: 输出目录
        train_ratio: 训练集比例
        val_ratio: 验证集比例
    """
    # 创建输出目录
    for split in ['images/train', 'images/val', 'images/test', 
                  'labels/train', 'labels/val', 'labels/test']:
        Path(f"{output_dir}/{split}").mkdir(parents=True, exist_ok=True)
    
    # 获取所有图片文件
    image_files = [f for f in os.listdir(image_dir) 
                   if f.endswith(('.jpg', '.jpeg', '.png'))]
    
    # 打乱顺序
    random.shuffle(image_files)
    
    # 计算分割点
    n_train = int(len(image_files) * train_ratio)
    n_val = int(len(image_files) * val_ratio)
    
    splits = {
        'train': image_files[:n_train],
        'val': image_files[n_train:n_train + n_val],
        'test': image_files[n_train + n_val:]
    }
    
    # 复制文件
    for split_name, files in splits.items():
        for filename in files:
            # 复制图片
            src_img = f"{image_dir}/{filename}"
            dst_img = f"{output_dir}/images/{split_name}/{filename}"
            shutil.copy(src_img, dst_img)
            
            # 复制标签(如果有)
            label_file = Path(image_dir).parent / 'labels' / f"{Path(filename).stem}.txt"
            if label_file.exists():
                dst_label = f"{output_dir}/labels/{split_name}/{Path(filename).stem}.txt"
                shutil.copy(label_file, dst_label)
    
    print(f"数据集划分完成!")
    print(f"  训练集: {len(splits['train'])} 张")
    print(f"  验证集: {len(splits['val'])} 张")
    print(f"  测试集: {len(splits['test'])} 张")

# 使用示例
split_dataset(
    image_dir='dataset/raw_images',
    output_dir='MaskDataSet',
    train_ratio=0.8,
    val_ratio=0.1
)

五、创建数据集配置文件

yolov5/data/ 目录下创建 data.yaml

可以参考yolov5\data\coco128.yaml来创建自己的data.yaml

复制代码
---
 **关键点**:
`names` 中的类别顺序必须与标注时的类别 ID 一致
 `nc` 必须等于类别总数
`path` 可以使用绝对路径或相对路径

六、模型选择与下载

6.1 YOLOv5 模型系列

模型 参数量 (M) FLOPs (B) mAP@50 速度 (V100) 适用场景
YOLOv5n 1.9 4.5 45.7 6.3ms 边缘设备、实时性要求极高
YOLOv5s 7.2 16.5 56.8 6.4ms 通用场景(推荐入门)
YOLOv5m 21.2 49.0 64.1 8.2ms 需要更高精度
YOLOv5l 46.5 109.1 67.3 10.1ms 高精度需求
YOLOv5x 86.7 205.7 68.9 12.1ms 最高精度

💡 新手推荐 :从 YOLOv5s 开始,显存需求较低(6GB+),训练速度快。

6.2 下载预训练权重

YOLOv5 会自动下载预训练权重,也可以手动下载:

bash 复制代码
# 方式一:训练时自动下载(首次训练会自动下载)
python train.py --data data/mask_data.yaml --weights yolov5s.pt

# 方式二:手动下载(更快)
# YOLOv5s
curl -L -o yolov5s.pt https://github.com/ultralytics/yolov5/releases/download/v7.0/yolov5s.pt

# YOLOv5m
curl -L -o yolov5m.pt https://github.com/ultralytics/yolov5/releases/download/v7.0/yolov5m.pt

# YOLOv5n(最小模型)
curl -L -o yolov5n.pt https://github.com/ultralytics/yolov5/releases/download/v7.0/yolov5n.pt

预训练权重下载位置:https://github.com/ultralytics/yolov5/releases


七、训练配置与运行

7.1 基础训练命令

bash 复制代码
python train.py \
    --data data/data.yaml \
    --weights yolov5s.pt \
    --epochs 100 \
    --batch-size 16 \
    --img 640 \
    --device 0

7.2 完整训练参数配置

根据实际硬件配置和需求调整 train.py 参数:

参数 默认值 说明 推荐值(6GB显存)
--data coco128.yaml 数据集配置文件路径 data/mask_data.yaml
--weights yolov5s.pt 预训练权重路径 yolov5s.pt
--cfg '' 模型架构文件 models/yolov5s.yaml
--epochs 100 训练轮数 100~300
--batch-size 16 批次大小 8~16(根据显存调整)
--img 640 输入图像尺寸 640
--device '' 计算设备 0(GPU0)或 cpu
--workers 8 数据加载线程数 4~8
--project runs/train 输出目录 自定义
--name exp 实验名称 mask_train
--hyp data/hyps/hyp.scratch-low.yaml 超参数文件 默认
--optimizer SGD 优化器 SGD / Adam / AdamW
--lr0 0.01 初始学习率 0.01
--momentum 0.937 SGD 动量 0.937
--weight-decay 0.0005 权重衰减 0.0005
--freeze [0] 冻结层数 可冻结 Backbone

7.3 不同显存大小的批次大小建议

GPU 显存 YOLOv5n YOLOv5s YOLOv5m YOLOv5l YOLOv5x
4GB 64 16 8 4 2
6GB 128 32 16 8 4
8GB 256 64 32 16 8
12GB+ 512 128 64 32 16

7.4 常用训练命令示例

bash 复制代码
# 示例1:使用 YOLOv5s 训练口罩数据集(推荐配置)
python train.py \
    --data data/data.yaml \
    --weights yolov5s.pt \
    --epochs 100 \
    --batch-size 16 \
    --img 640 \
    --device 0 \
    --name mask_yolov5s \
    --project runs/train

# 示例2:从头训练(不使用预训练权重 不太建议)
python train.py \
    --data data/data.yaml \
    --weights '' \
    --cfg models/yolov5s.yaml \
    --epochs 200 \
    --batch-size 16 \
    --img 640

# 示例3:多尺度训练(提升小目标检测)
python train.py \
    --data data/data.yaml \
    --weights yolov5s.pt \
    --epochs 100 \
    --batch-size 8 \
    --img 640 \
    --multi-scale

# 示例4:冻结 Backbone 训练(加速收敛)
python train.py \
    --data data/data.yaml \
    --weights yolov5s.pt \
    --epochs 50 \
    --batch-size 32 \
    --img 640 \
    --freeze 10

7.5 修改模型配置文件(如需调整类别数)

如果修改了模型结构,编辑 models/yolov5s.yaml

yaml 复制代码
# YOLOv5s 模型配置(仅列出最后一层需要修改的地方)
# ...
# yolov5s.yaml 中的检测头部分

# 检测头(根据 nc 修改)
nc: 2  # 修改为你的类别数(这里口罩数据集有 2 类)

八、训练过程监控

8.1 训练输出目录

训练完成后,结果保存在 runs/train/ 目录下:

复制代码
runs/train/exp/
├── weights/
│   ├── best.pt       # 验证集上 mAP 最高的模型
│   └── last.pt       # 最后一轮的模型
├── results.csv       # 训练指标记录(可用 Excel 打开)
├── results.png       # 训练曲线图
├── args.yaml         # 训练参数配置
└── ...

8.2 训练曲线解读

训练完成后会生成 results.png,包含以下曲线:

曲线 含义 理想状态
box_loss 边界框定位损失 持续下降
obj_loss 目标置信度损失 持续下降
cls_loss 分类损失 持续下降
precision 精确率 持续上升趋于稳定
recall 召回率 持续上升趋于稳定
mAP@0.5 IoU=0.5 时的 mAP 持续上升趋于稳定
mAP@0.5:0.95 多 IoU 阈值平均 mAP 持续上升趋于稳定

图2:训练过程曲线,监控各项指标的变化

8.3 TensorBoard 可视化

bash 复制代码
# 安装 tensorboard
pip install tensorboard

# 启动 tensorboard
tensorboard --logdir runs/train

# 浏览器打开 http://localhost:6006

九、模型推理与检测

9.1 使用 detect.py 推理

bash 复制代码
# 检测单张图片
python detect.py \
    --weights runs/train/exp/weights/best.pt \
    --source data/images/test/img001.jpg \
    --device 0

# 检测多张图片
python detect.py \
    --weights runs/train/exp/weights/best.pt \ 
    --source data/images/test/ \
    --device 0 \
    --save-img

# 检测视频
python detect.py \
    --weights runs/train/exp/weights/best.pt \
    --source video.mp4 \
    --device 0 \
    --save-video

# 摄像头实时检测
python detect.py \
    --weights runs/train/exp/weights/best.pt \
    --source 0 \
    --device 0

# 批量检测并保存结果
python detect.py \
    --weights runs/train/exp/weights/best.pt \ 
    --source data/images/val/ \
    --device 0 \
    --project runs/detect \
    --name exp \
    --save-txt \
    --save-conf

检测结果保存在 runs/detect/ 目录。

9.2 Python API 推理

python 复制代码
import torch
from pathlib import Path

# 方法一:使用 PyTorch Hub(最简洁)
model = torch.hub.load(
    'ultralytics/yolov5', 
    'custom', 
    path='runs/train/mask_yolov5s/weights/best.pt'
)

# 设置置信度阈值
model.conf = 0.5  # 置信度阈值
model.iou = 0.45  # NMS IoU 阈值

# 推理
img = 'data/images/test/img001.jpg'
results = model(img)

# 显示结果
results.show()

# 保存结果
results.save()

# 获取检测结果(DataFrame 格式)
df = results.pandas().xyxy[0]
print(df)
# 输出示例:
#     xmin    ymin    xmax    ymax  confidence  class    name
# 0  100.5    80.3    250.4   300.2      0.92     0   with_mask
# 1  300.1    120.5   400.2   350.1      0.87     1   without_mask

# 方法二:使用 YOLOv5 模块
from models.common import DetectMultiBackend
from utils.general import non_max_suppression

# 加载模型
model = DetectMultiBackend('runs/train/mask_yolov5s/weights/best.pt')
img = torch.zeros((1, 3, 640, 640))  # 输入图像

# 推理
pred = model(img)
pred = non_max_suppression(pred)

print(f"检测到 {len(pred[0])} 个目标")

十、模型验证

10.1 在验证集上评估模型

bash 复制代码
python val.py \
    --data data/data.yaml \
    --weights runs/train/exp/weights/best.pt \
    --img 640 \
    --batch-size 32 \
    --device 0

10.2 输出指标

验证完成后会输出以下指标:

指标 含义
mAP@0.5 IoU=0.5 时的平均精度
mAP@0.5:0.95 多 IoU 阈值的平均精度
Precision 精确率
Recall 召回率
FPS 推理速度(帧/秒)

十一、模型导出与部署

11.1 导出为不同格式

bash 复制代码
# 导出为 ONNX(通用格式,推荐)
python export.py \
    --weights runs/train/exp/weights/best.pt \
    --include onnx \
    --img 640

# 导出为 TensorRT(NVIDIA GPU 加速)
python export.py \
    --weights runs/train/exp/weights/best.pt \
    --include engine \
    --device 0 \
    --half

# 导出为 TFLite(移动端/嵌入式)
python export.py \
    --weights runs/train/exp/weights/best.pt \
    --include tflite \
    --img 640

# 导出为 CoreML(iOS)
python export.py \
    --weights runs/train/exp/weights/best.pt \
    --include coreml \
    --img 640

11.2 导出格式对比

格式 适用场景 优点 缺点
ONNX 跨平台部署 通用性强,工具链完善 推理速度一般
TensorRT NVIDIA GPU 速度最快(3-5x 加速) 仅支持 NVIDIA
TFLite 移动端/嵌入式 Android/iOS 友好 精度略有损失
CoreML iOS/macOS Apple 生态原生支持 仅支持 Apple 设备

十二、常见问题与解决方案

Q1:训练时显存不足(OOM)

bash 复制代码
# 解决方案:减小批次大小
python train.py \
    --data data/data.yaml \
    --weights yolov5s.pt \
    --batch-size 8      # 减小批次
    --img 416            # 减小输入尺寸

Q2:训练中断,想恢复训练

bash 复制代码
python train.py \
    --data data/data.yaml \
    --weights runs/train/exp/weights/last.pt \
    --resume runs/train/exp/weights/last.pt \
    --epochs 200

Q3:GPU 利用率低,训练慢

bash 复制代码
# 解决方案:增加数据加载线程数
python train.py \
    --data data/data.yaml \
    --weights yolov5s.pt \
    --workers 8 \
    --batch-size 16

Q4:检测不到小目标

bash 复制代码
# 解决方案:使用更高分辨率 + 多尺度训练
python train.py \
    --data data/data.yaml \
    --weights yolov5s.pt \
    --img 1280 \
    --batch-size 8 \
    --multi-scale

Q5:模型过拟合

bash 复制代码
# 解决方案:使用数据增强 + 正则化
python train.py \
    --data data/data.yaml \
    --weights yolov5s.pt \
    --epochs 100 \
    --hyp data/hyps/hyp.scratch.yaml \
    --augment

十三、完整工作流程总结

复制代码
1. 环境准备
   ├── 安装 Python >= 3.8
   ├── 安装 PyTorch + CUDA
   └── 安装 yolov5 依赖
           ↓
2. 代码拉取
   └── git clone https://github.com/ultralytics/yolov5
           ↓
3. 数据集准备(MaskDataSet)
   ├── 收集图片
   ├── 使用 LabelImg 标注
   ├── 划分训练/验证/测试集
   └── 生成 YOLO 格式标签
           ↓
4. 创建数据集配置
   └── 编辑 data/mask_data.yaml
           ↓
5. 模型训练
   ├── 选择模型(yolov5s.pt)
   ├── 配置训练参数
   └── python train.py --data data/mask_data.yaml --weights yolov5s.pt
           ↓
6. 训练监控
   ├── 观察 loss 曲线
   ├── 观察 mAP 指标
   └── 使用 TensorBoard 可视化
           ↓
7. 模型推理
   ├── python detect.py
   └── Python API 推理
           ↓
8. 模型部署
   ├── export.py 导出模型
   └── 部署到目标平台

十四、小结

本文详细介绍了 YOLOv5 的完整使用流程:

  1. 环境配置:PyTorch + CUDA 环境搭建
  2. 代码拉取:从 GitHub 克隆 YOLOv5 仓库
  3. 数据集准备:MaskDataSet 的 YOLO 格式标注与划分
  4. 配置文件 :创建 mask_data.yaml 数据集配置
  5. 模型训练:完整的训练命令与参数调优
  6. 推理检测:使用 detect.py 和 Python API 进行目标检测
  7. 模型导出:导出为 ONNX/TensorRT/TFLite 等格式

YOLOv5 的工程化设计使得从研究到生产的路径非常顺畅,熟练掌握这套流程后,可以快速将目标检测技术应用到任何实际项目中。


参考资源

  1. YOLOv5 GitHub 仓库:https://github.com/ultralytics/yolov5
  2. YOLOv5 官方文档:https://docs.ultralytics.com
  3. YOLOv5 训练自定义数据教程:https://docs.ultralytics.com/yolov5/tutorials/train_custom_data/
  4. 预训练权重下载:https://github.com/ultralytics/yolov5/releases
相关推荐
guo_xiao_xiao_5 小时前
YOLOv11城市道路自行车目标检测数据集-999张-Bicycle-1
人工智能·yolo·目标检测
guo_xiao_xiao_5 小时前
YOLOv11机场跑道大型客机-飞机目标检测数据集-201张-Airplane-1_2_2
人工智能·yolo·目标检测
海域云-罗鹏5 小时前
豆包开启付费订阅,想白嫖越来越难了,企业不如部署自己的算力服务器
服务器·人工智能·github
遁神银灵子5 小时前
OMO(oh-my-openagents)插件在OpenCode Desktop v1.4.33以上版本失效问题研究
github
羊羊小栈6 小时前
基于「YOLO目标检测 + 多模态AI分析」的宠物猫狗健康智能检测分析预警系统
人工智能·yolo·目标检测·计算机视觉·毕业设计·大作业
码流怪侠7 小时前
【GitHub】TextGen:开源本地大模型运行平台的终极解决方案
python·程序员·github
小雨青年7 小时前
GitHub Copilot Commit Message 生成与自定义配置优化指南
人工智能·github·copilot
无限进步_7 小时前
【C++】AVL树完全解析:从平衡因子到四种旋转
c语言·开发语言·数据结构·c++·后端·算法·github
南斯拉夫的铁托7 小时前
YOLO学习笔记
笔记·学习·yolo