pytorch中的DataLoader

最新推荐文章于 2025-05-14 23:05:40 发布

原创

最新推荐文章于 2025-05-14 23:05:40 发布 · 1.6k 阅读

17 ·

CC 4.0 BY-SA版权

文章标签：

#pytorch #人工智能 #数据加载

通常在训练时我们会将数据集分成若干小的、随机的批（batch），这个操作当然可以手动操作，但是pytorch里面为我们提供了API让我们方便地从dataset中获得batch，DataLoader就是来解决这个问题的。

DataLoader的本质是一个可迭代对象，即经过DataLoader的返回值为一个可迭代的对象，一般的操作是：1、创建一个 dataset 对象；2、创建一个DataLoader对象；3、遍历这个DataLoader对象，将data, label加载到模型中进行训练。

torch.utils.data.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, multiprocessing_context=None)

dataset：官方文档给出的解释是 “ 从中加载数据的数据集 ” 。
batch_size（int，可选）：每个批次要加载的样本数（默认值：1）。
shuffle（bool，可选）：设置为“True”可在每个epoch重新排列数据（默认值：“False”）。一个epoch表示数据集的数据全部使用一遍。
sampler（sampler，可选）：定义从数据集中提取样本的策略。如果指定，“shuffle”必须为False。batch_sampler 类似，表示一次返回一个batch的索引。
num_workers（int，可选）：用于数据加载的子进程数，0 表示将在主进程中加载数据。（默认值：0）。换句话说，num_workers = 0 表示在主进程中加载数据而不使用任何额外的子进程；若大于0，表示开启多个进程。进程越多，处理数据的速度越快，但会使电脑性能下降，占用更多的内存。
collate_fn（可调用，可选）：表示合并样本列表以形成小批量的Tensor对象。
pin_memory（bool，可选）：表示要将load进来的数据是否要拷贝到pin_memory区中，其表示生成的Tensor数据是属于内存中的锁页内存区，这样将Tensor数据转义到GPU中速度就会快一些，默认为False。如果为“True”，数据加载程序将在返回张量之前将张量复制到CUDA固定内存中。通常情况下，数据在内存中要么以锁页的方式存在，要么保存在虚拟内存(磁盘)中，pin_memory设置为True后，数据直接保存在锁页内存中，后续直接传入CUDA；否则需要先从虚拟内存中传入锁页内存中，再传入CUDA，这样就比较耗时了。
drop_last（bool，可选）：当整个数据长度不能够整除batch_size，选择是否要丢弃最后一个不完整的batch，默认为False。设置为“True”时可以删除最后一个不完整的批次（batch）。

1 处理数据集

在许多情况下，在众所周知的数据集（如 MNIST 或 CIFAR）上训练神经网络，可以实现预测的准确率超过 90%。原因是，这些数据集组织整齐且易于预处理。但是，当处理自己的数据集时，要实现高精度非常棘手且具有挑战性。

我们将快速浏览一下 PyTorch 库中包含的数据集。PyTorch 带有几个内置的数据集，所有这些都预加载在torch.datasets类中。torch有torchvision，该torch包实现神经网络所需的所有核心类和方法，torchvision包含流行的数据集、模型架构和计算机视觉的常见图像转换。