文章目录
0. 前言
按照国际惯例,首先声明:本文只是我自己学习的理解,虽然参考了他人的宝贵见解及成果,但是内容可能存在不准确的地方。如果发现文中错误,希望批评指正,共同进步。
在深度学习中,数据的预处理和加载方式对模型训练的效率与效果具有重要影响。PyTorch提供了一种强大的工具——DataLoader
,它能够高效地将数据集转化为适合模型训练的小批量数据,并支持多线程并行加载机制,极大地提升了数据读取速度。本文将详细介绍PyTorch中的DataLoader。
本文的说明参考了PyTorch官网文件:https://pytorch.org/docs/stable/data.html
DataLoader通常会结合ImageFolder和transforms类(即构建Dataset过程)一起使用,这两个类已经在此前文章中专题说明过:
1. DataLoader的功能
根据DataLoader的官方文档说明,将从以下5个方面说明DataLoader的功能:
1.1 可处理映射式/可迭代式数据集
PyTorch 的 DataLoader
能够处理两种形式的数据集:映射式数据集(map-style)和可迭代式数据集(iterable-style)。映射式数据集指的是那些可以通过索引直接访问其元素的数据集,它们需要实现 __getitem__
方法和可选的 __len__
方法。例如,torch.utils.data.Dataset
子类就是这样一种数据集,可以通过 dataset[i]
获取第 i
个样本。
可迭代式数据集则不依赖于索引访问,而是可以直接迭代出数据,这类似于 Python 中的迭代器协议。DataLoader
能够兼容这两种风格的数据集,并将它们转化为合适的形式以供模型训练使用。
1.2 可自定义数据加载顺序
DataLoader
支持通过 sampler
参数来自定义数据加载的顺序。默认情况下,如果设置了 shuffle=True
,那么在每个 epoch 开始时,数据加载器会打乱数据集的顺序。但也可以传入自定义的采样器,如 RandomSampler
实现随机抽样,或 SequentialSampler
保持原有顺序,甚至可以使用 WeightedRandomSampler
来实现加权随机抽样等复杂逻辑。
1.3 可自动批量化打包数据
DataLoader
自动将数据集中的样本打包成小批量,这是通过设置 batch_size
参数来实现的。每次调用 DataLoader
的迭代器时,都会返回一个包含 batch_size
个样本的数据批次,这对于训练深度学习模型是非常关键的,因为大多数模型都需要按照批次进行前向传播和反向传播计算。
1.4 可支持多进程加载
为了加速数据加载过程,特别是对于大型数据集,DataLoader
提供了多进程支持。通过设置 num_workers
参数,可以启动多个工作进程来并发地加载数据。这意味着数据准备和模型训练可以同时进行,极大地提高了整体效率。需要注意的是,启用多进程数据加载时需要考虑数据集是否线程安全,并确保在CPU或系统资源充足的环境下运行。
注意:num_workers
参数并不一定是越大越好。以下是考虑 num_workers
设置时需要权衡的因素:
-
系统资源:
- CPU核心数:
num_workers
应设置得小于或等于可用 CPU 核心数(包括超线程)。若设置得过高,可能会导致过多的上下文切换,反而降低性能。 - 内存限制:增加
num_workers
可能会增加内存消耗,因为每个工作进程都会缓存一部分数据。过高的num_workers
可能会导致内存溢出。
- CPU核心数:
-
I/O 瓶颈:
- 如果数据读取主要是由硬盘 I/O 速率决定的瓶颈,那么超过某个点后增加
num_workers
不会带来进一步的速度提升,反而可能由于争抢 I/O 资源而造成负面影响。
- 如果数据读取主要是由硬盘 I/O 速率决定的瓶颈,那么超过某个点后增加
-
GPU 同步:
- 当数据加载速度远大于 GPU 计算速度时,更多的
num_workers
可能不会显著提高训练效率,因为 GPU 处理速度成为瓶颈。
- 当数据加载速度远大于 GPU 计算速度时,更多的
-
同步与异步行为:
- PyTorch 的 DataLoader 默认实现了一个队列系统来进行数据加载的同步操作。过多的
num_workers
可能会导致队列中积累过多的数据,这些数据在被 GPU 使用前需要等待,因此并不会提高整体吞吐量。
- PyTorch 的 DataLoader 默认实现了一个队列系统来进行数据加载的同步操作。过多的
通常的经验值可能是把num_workers
设定在 CPU 核心数的一半到全部之间。不过最佳实践是要根据具体的硬件配置、数据集大小和读取速度以及模型训练速度等因素进行调整。在某些情况下,如遇到操作系统兼容性问题或者为了避免不必要的复杂性,也可能需要将 num_workers
设置为较小的值甚至 0。在 Windows 系统中,由于进程间通信的限制,有时必须将 num_workers
设置为 0 才能避免错误。
1.5 可pin住内存
如果是在 GPU 上进行训练,DataLoader
可以通过设置 pin_memory=True
来自动将 CPU 内存中的数据拷贝至 CUDA 可以直接访问的内存区域(即“pin”住内存),这样在数据从 CPU 到 GPU 的转移过程中可以享受到更快的速度。这是因为被pin住内存可以利用异步内存复制操作,避免同步等待,从而使得数据流水线更为顺畅。
2. DataLoader的调用
在PyTorch中,DataLoader是基于torch.utils.data
接口进行工作的,DataLoader也是torch.utils.data
中的核心类。
At the heart of PyTorch data loading utility is the torch.utils.data.DataLoader class.
2.1 DataLoader的调用方法
DataLoader的调用方法如下:
from torch.utils.data import DataLoader
dataset = ...
loader = DataLoader(dataset, batch_size=1, shuffle=False, sampler=None,
batch_sampler=None, num_workers=0, collate_fn=None,
pin_memory=False, drop_last=False, timeout=0,
worker_init_fn=None, prefetch_factor=2,
persistent_workers=False)
2.2 DataLoader的参数说明
-
dataset: (Dataset) 必须是一个实现了
__getitem__
和__len__
方法的数据集对象,用于定义如何访问和获取样本及其对应的标签。 -
batch_size: (int) 指定每个批次加载多少个样本。这是训练神经网络时常见的参数,控制着每一步梯度更新所使用的样本数量。
-
shuffle: (bool) 若为
True
,则在每个 epoch 开始时都会对数据集进行随机打乱顺序。这对于防止模型过拟合和确保模型看到所有样本组合至关重要,尤其是在训练阶段。 -
sampler: (Sampler) 可选参数,用于指定数据抽样的策略。如果提供了
sampler
,则shuffle
参数将被忽略。例如,可以使用此参数实现分布式训练时的数据分片。 -
batch_sampler: (Sampler) 可直接指定一个批量抽样器,它可以返回一批批索引而不是单个索引。如果指定了
batch_sampler
,则batch_size
和shuffle
将被忽略。 -
num_workers: (int) 上面已经详细说明,不再赘述。
-
collate_fn: (callable) 自定义函数,用于合并多个样本到一个批次。默认的
collate_fn
会堆叠具有相同形状的张量。用户可以自定义此函数以满足特殊的数据整合需求。 -
pin_memory: (bool) 若为
True
,则在数据加载后将其移到 CUDA 可以直接访问的页锁定内存中,从而加快数据从 CPU 到 GPU 的传输速度。 -
drop_last: (bool) 若为
True
,并且最后一个批次的样本数量小于batch_size
,则丢弃该批次。这在保证所有批次样本数量一致时很有用。 -
timeout: (float) 设置数据加载过程中阻塞的最大秒数。如果设置为0,则无限期等待。
-
worker_init_fn: (callable) 用户自定义的回调函数,用于初始化每个工作进程。可以在每个工作进程中设置不同的随机种子等。
-
prefetch_factor: (int) 预取因子,决定了工作进程在向主进程提交批次的同时,提前生成多少个额外的批次。增加此值可以减少潜在的 I/O 瓶颈,但也可能增加内存占用。
-
persistent_workers: (bool) 若为
True
,则保留工作进程在多个数据加载迭代之间,这样可以避免每次重新启动工作进程带来的开销,尤其在长时间运行的任务中效果更明显。然而,这要求num_workers > 0
并且multiprocessing.get_start_method()
返回 ‘fork’ 或 ‘spawn’。
3. DataLoader的使用实例
创建了一个 ImageFolder
数据集并应用了 transforms.ToTensor()
转换之后,你已经正确地设置了数据预处理流程,确保了 DataLoader
在批处理时接收到的是 Tensor 类型的数据。接下来要查看 DataLoader
中的特定元素,你可以通过迭代的方式来访问它们:
from torchvision import transforms
from torchvision.datasets import ImageFolder
from torch.utils.data import DataLoader
# 数据预处理
transform = transforms.Compose([transforms.Resize((200,200)), transforms.ToTensor()])
# 加载数据集
dataset = ImageFolder('D:\\DL\\pretrain\\hymenoptera\\hymenoptera_data', transform=transform)
# 创建 DataLoader
loader = DataLoader(dataset,batch_size=4)
# 通过迭代的方式访问 DataLoader 中的元素
for i, (images, labels) in enumerate(loader):
if i == 0: # 仅显示第一个批次的数据
print(f"第{i}个批次的图像张量:")
print(images.shape) # 显示图像张量的形状
print("对应的标签:", labels)
break
上述代码会打印出 DataLoader 中的第一个批次(索引为0)的图像张量及其对应的类别标签。输出为:
第0个批次的图像张量:
torch.Size([4, 3, 200, 200])
对应的标签: tensor([0, 0, 0, 0])
如果你想查看特定索引位置的样本(即上文所述的映射式map-style),但不是按批次而是直接查看原始数据集中的某个样本,那么可以直接从 ImageFolder
数据集中索引该样本,而非通过 DataLoader
:
# 获取数据集中的特定样本(假设索引为10)
sample_idx = 10
image, label = dataset[sample_idx]
print(f"索引 {sample_idx} 的图像张量:")
print(image.shape)
print("对应的标签:", label)
输出为:
索引 10 的图像张量:
torch.Size([3, 200, 200])
对应的标签: 0
再次强调,在实际应用中,由于 DataLoader
主要是用来进行批处理的,所以直接从其中索引单个元素并不常见。如需查看单个样本,通常是从原始 dataset
中访问。如果确实需要从 DataLoader 中一次性获取单个样本而不是一批次,需要特殊处理,例如通过 next(iter(loader))
或者额外编写逻辑来实现。
4. 总结
最后我想再总结下 DataLoader
和 Dataset
的关系:
DataLoader
依赖于Dataset
来获取原始数据,它的目的是为了更好地管理和高效地喂入数据给训练过程。- 使用时,首先需要基于
Dataset
构建好数据集实例,然后将这个数据集实例传给DataLoader
构造函数,配置好加载参数后得到一个数据加载器。
总的来说,Dataset
负责定义数据源和访问逻辑,而 DataLoader
负责根据这些定义好的逻辑,按需以适合训练的形式加载和提供数据。