数据加载

相关源文件

• benchmarks.py
• utils/augmentations.py
• utils/autoanchor.py
• utils/dataloaders.py
• utils/segment/augmentations.py
• utils/segment/dataloaders.py

本文档解释了YOLOv5中的数据加载流程和机制。它涵盖了数据如何被加载、处理和准备用于训练和推理操作。有关数据集配置和格式的信息，请参阅数据集。有关数据增强技术的详细信息，请参阅数据增强。

概述

YOLOv5的数据加载系统旨在从各种来源高效地加载图像、视频及其相应的标签。它包括批量处理、缓存、预处理和应用基本转换的功能，以准备用于模型训练或推理的数据。

该系统优先考虑

• 通过缓存和多线程提高性能
• 处理不同数据格式和来源的灵活性
• 支持分布式训练
• 高效的内存管理

来源: utils/dataloaders.py1-68

数据加载架构

数据加载流程

来源: utils/dataloaders.py160-218 utils/segment/dataloaders.py21-83

类层次结构

来源: utils/dataloaders.py322-430 utils/dataloaders.py433-526 utils/dataloaders.py264-320 utils/dataloaders.py537-1038 utils/dataloaders.py264-320

创建数据加载器

YOLOv5 使用 `create_dataloader()` 函数创建数据加载器，该函数负责初始化相应的数据集类并将 PyTorch DataLoader 配置为正确的设置。

主要参数

参数	描述
path	图像或数据集 YAML 文件的路径
imgsz	图像尺寸（高=宽）
batch_size	加载的批量大小
stride	模型步长（用于矩形训练）
single_cls	视为单类数据集
hyp	用于数据增强的超参数
augment	应用数据增强
cache	将图像缓存到 RAM 或磁盘
rect	使用矩形训练
workers	工作线程数

来源: utils/dataloaders.py160-218

创建过程

来源: utils/dataloaders.py160-218 utils/segment/dataloaders.py21-83

核心数据集类

LoadImagesAndLabels

这是用于训练和验证的主要数据集类。它负责加载图像及其关联的标签。

关键功能：

• 从磁盘加载图像和标签
• 缓存以提高性能
• 矩形批量处理以实现高效训练
• 马赛克数据增强支持
• 支持图像权重

来源: utils/dataloaders.py537-1038

LoadImages

用于对图像和视频进行推理。处理图像文件和视频文件，逐帧处理。

关键功能：

• 支持多种图像格式
• 使用OpenCV进行视频处理
• 图像缩放以获得正确尺寸
• 视频逐帧迭代

来源: utils/dataloaders.py322-430

LoadStreams

专门用于视频流（包括网络摄像头、RTSP 流和 YouTube 视频）的实时推理的类。

关键功能：

• 多线程流读取
• 支持 YouTube 链接
• 实时处理多个流
• 流失败时自动重新连接

来源: utils/dataloaders.py433-526

LoadScreenshots

加载和处理指定屏幕区域的屏幕截图，用于实时检测。

关键功能：

• 使用MSS库进行屏幕捕获
• 可配置的屏幕区域
• 逐帧迭代

来源: utils/dataloaders.py264-320

训练数据加载流程

训练数据加载流程涉及从文件路径到批量张量的多个步骤

来源: utils/dataloaders.py537-1038 utils/dataloaders.py160-218

缓存系统

YOLOv5 实现了一个复杂的缓存系统来加速训练

标签缓存

每个图像的标签都缓存到一个 `.cache` 文件中，以避免重复解析标注文件。缓存包含：

• 标签数据（类别、边界框坐标）
• 图像形状
• 分割数据（如果适用）
• 数据集文件的哈希值，用于检测更改

来源: utils/dataloaders.py722-760

图像缓存

图像可以缓存到 RAM 或磁盘上，格式为 `.npy` 文件。

• RAM 缓存：将所有图像加载到内存中以实现最快的访问
• 磁盘缓存：将图像转换为 NumPy 数组并存储在磁盘上
• 可用内存检查可防止 OOM 错误

来源: utils/dataloaders.py683-702 utils/dataloaders.py703-720 utils/dataloaders.py848-872 utils/dataloaders.py873-878

特殊加载技术

矩形训练

矩形训练将具有相似纵横比的图像分组到批次中，以最小化填充并最大化 GPU 利用率。

进程

1. 按纵横比对图像进行排序
2. 按相似的纵横比分组到批次中
3. 计算最优批次形状
4. 调整填充以保持步长兼容性

来源: utils/dataloaders.py657-680

马赛克加载

马赛克加载将四张训练图像合并成一张，增加了单张训练样本中物体和上下文的多样性。

进程

1. 选择一张目标图像和三张随机图像
2. 将它们放在 2x2 网格中，并随机偏移
3. 调整标签和分割以匹配新位置
4. 应用额外的增强（随机透视等）

来源: utils/dataloaders.py879-959

分布式训练支持

YOLOv5 通过以下方式提供对分布式训练的支持

SmartDistributedSampler

扩展了 PyTorch 的 DistributedSampler 以确保

• 基于 epoch 和 seed 的确定性打乱
• 跨 GPU 的平衡数据分布
• 与矩形训练兼容

来源： utils/dataloaders.py132-157

Torch Distributed Zero First

在分布式环境中，该上下文管理器确保数据集缓存操作仅由主进程执行。

来源： utils/dataloaders.py183-184

批处理

InfiniteDataLoader

一个自定义数据加载器，通过重用工作进程来避免每个 epoch 重新创建它们带来的开销。

主要功能

• 连续数据迭代，无工作进程重启
• 兼容标准的 PyTorch DataLoader API
• 内部使用重复采样器

来源： utils/dataloaders.py221-245 utils/dataloaders.py246-262

Collate Functions (批处理函数)

自定义的批处理函数，用于将单个数据集项组合成批次

• collate_fn：标准的目标检测任务批处理函数
• collate_fn4：用于马赛克训练，将四张图像合并在一起的批处理函数
• 用于处理掩码的特定于分割的版本

来源： utils/dataloaders.py215-216 utils/segment/dataloaders.py297-304

图像处理和路径处理

通用图像操作

• letterbox：调整图像大小并填充，使其适应目标尺寸，同时保持宽高比
• exif_size, exif_transpose：处理相机 EXIF 方向数据
• img2label_paths：从图像路径生成标签文件路径

来源： utils/augmentations.py122-152 utils/dataloaders.py83-117 utils/dataloaders.py529-535

分割任务数据加载

对于分割任务，LoadImagesAndLabelsAndMasks 类扩展了 LoadImagesAndLabels 来处理额外的掩码数据

主要新增功能

• 掩码降采样以提高效率
• 多边形到掩码的转换
• 实例分割的重叠处理
• 保持掩码对齐的特殊增强

来源： utils/segment/dataloaders.py86-236 utils/segment/dataloaders.py307-366

最佳实践

• 对于大型数据集，请使用磁盘缓存（cache='disk'），而不是内存缓存
• 在不使用马赛克增强时，启用矩形训练（rect=True）以提高效率
• 根据您的 CPU 能力调整工作进程数（通常为 2-8）
• 对于分布式训练，请确保图像在 GPU 之间均匀分布
• 在设置批次大小和图像分辨率时，请考虑内存影响