torch.utils.data学习笔记

其实就是翻译了一下，
官网：https://pytorch.org/docs/stable/data.html#torch.utils.data.DataLoader

在 Pytorch 的数据加载工具中，torch.utils.data.DataLoader 类起到核心的作用，它是在数据集上的一个 Python 迭代器，并支持以下内容：

Map 与迭代器类型的数据集；
自定义的数据加载指令；
自动分批；
单进程与多进程的数据加载；
自动内存pinning

这些选项通过 DataLoader 的构造器的参数进行配置，该类的构造器记为：

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)

下列章节将详细的描述这些选项的效果与使用方法。

1. 数据集类型

在 DataLoader构造器参数中最重要的一项是 Dataset，它指定了用于加载数据的数据集实例。Pytorch 中提供了两种不同的数据集类型：

Map 类型的数据集
迭代器类型的数据集

1.1. Map 类型的数据集

Map 类型的数据集是指实现了 __getitem__() 与 __len__() 的一个类，该类是一个由索引或者键值到数据样本的一个映射。
例如，dataset[idx] 表示从磁盘的文件夹中读取第 i 张图像及相关的标签。
详情参考 Dataset 。

1.2. 迭代器类型的数据集

迭代器类型的数据集是 IterableDataset 类的子集，它实现了函数接口 __iter__()，表示在数据样本上的迭代器。对于那些随机读写比较困难甚至不大可能发生的情况，以及批次大小依赖于待取数据的情况，该类型的数据集是比较合适的。
例如，iter(dataset) 可以返回一个从远程数据库甚至是实时日志中读取的数据流。
详情参考 IterableDataset 。

note:
当通过多进程数据加载方案使用 IterableDataset 类时，对于同一份数据，不同的工作进程将各自产生一个相同的复制品，这些数据必须予以不同的方式进行安排，以消除重复数据。具体的操作方式可以参考 IterableDataset 的文档。

2. 数据加载指令与采样工具

对于迭代器类型的数据集来说，数据加载指令完全是由用户定义的迭代器决定的。它可以使读取大块数据或者动态设定批次大小的实现方式变的很简洁（比如，每次可以选定（yield）一个批次的样本）。
本节余下的内容将主要讨论Map 类型的数据集。torch.utils.data.Sampler类用于指定数据加载过程中使用的索引或者键值序列。它表示由索引到数据集的可迭代对象。比如，在使用SGD的时，Sampler 可以将索引随机排列，每次选定（yield）一个样本，或者为 mini-batch SGD 选定（yield）一小组样本。
基于DataLoader 参数中的 shuffle 值会自动创建一个顺序或者乱序的采样工具。此外，用户也可以通过参数 sampler 里指定一个自定义的 Sampler 实例，以在每次采样中获取下一个索引或者键。
batch_sampler 参数可以用来设置自定义的每次能选定（yield）一组索引值的采样工具。自动分批次的操作可以通过参数 batch_size 与 drop_last 设置，详情请参考下一节。

note:
由于迭代器类型的数据集中不一定有索引与键，所以参数 batch_size 与 drop_last 无法用于迭代器类型的数据集。

2.1 加载分批次或者不分批次的数据

DataLoader 通过参数 batch_size、drop_last 与 batch_sampler 自动地将采样数据整合进一个批次。

2.1.1. 自动分批（默认）

这一项最常用，其相当于取一些最小批次数据并将它们整合进同一个批次的样本中。也就是说在张量中有一个维度是批维度（通常是第一项）。
当参数 batch_size （默认为1）不为 None 时，数据加载工具返回的是分批的数据，而不是独立数据。参数 batch_size 与 drop_last 用于指定数据加载工具获取每批数据键的方式。对 Map 类型的数据来说，用户也可以通过设置 batch_sampler，每次返回一个键的列表。

note
参数 batch_size 与 drop_last 其实是在根据参数 sampler 构建 batch_sampler。对 Map 类型的数据集来说，sampler 也可以通过用户设置，或者基于参数 shuffle 构建；对迭代器类型的数据集来说，sampler 是一个虚拟的无限采样器。关于采样器的详情，可参考这里。

note
当以多进程的方式从迭代器类型的数据集中取数据时，在每个进程中，如果最后一个批次中的数据数量少于参数 batch_size 中设置的值，参数 drop_last 会将该批次舍弃。

在通过采样工具获取到一个样本列表之后，通过参数 collate_fn 设置的函数负责把列表整理进数据批。
设置参数 collate_fn 之后，从 Map 类型的数据集中加载数据的方式等价于：

for indices in batch_sampler:
    yield collate_fn([dataset[i] for i in indices])

从迭代器类型的数据集中加载数据的方式等价于：

dataset_iter = iter(dataset)
for indices in batch_sampler:
    yield collate_fn([next(dataset_iter) for _ in indices])

传入参数 collate_fn 的函数可以是自定义的。比如将序列数据填充到批次的最大长度。关于参数 collate_fn 的详细信息可参考这里。

2.1.2. 取消自动分批

有时，用户可能想在数据集代码中人工控制分批操作，或者仅仅想加载孤立样本。例如，直接加载分批的数据可能更加便捷（比如从数据库中大批的读取数据或者从内存中读取连续数据），或者批次的尺寸是依赖于数据的，又或者程序是在孤立样本上进行工作的。在这些场景中，自动分批可能不再适用（它是使用 collate_fn 整理样本），由数据加载工具直接返回数据集实例的每一项的操作可能会更方便。
当 batch_size 与 batch_sampler 同时设置为 None 时（batch_sampler 默认为 None），自动分批会被关闭。从数据集中获取的样本都会由参数 collate_fn 设置的函数进行处理。
自动分批功能被关闭之后，默认的 collate_fn 仅仅会将 NumPy 数组转换为 PyTorch 中 Tensor 的类型，其他不变。
此时，从 Map 类型的数据集中加载数据的方式等价于：

for index in sampler:
    yield collate_fn(dataset[index])

从迭代器类型的数据集中加载数据的方式等价于：

for data in iter(dataset):
    yield collate_fn(data)

关于参数 collate_fn 的详细信息可参考这里。

3. `collate_fn` 的使用

在自动分批开启与关闭的情况下，collate_fn 的作用是不相同的。
自动分批关闭时，collate_fn 的函数被孤立数据样本调用，其结果通过数据加载迭代器返回。
自动分批开启时，collate_fn 的函数被数据样本的列表调用，它将输入的数据样本整合进一批中，以便于从数据加载迭代器中返回。接下来讨论在自动分批开启时默认的 collate_fn 的表现。
以一个例子说明，如果每个数据样本由3通道的图像以及一个整型标签组成，亦即数据集的每一项返回一个元组(图像，标签)，默认的 collate_fn 将这些元组的列表整合为一个元组，该元组由分批后的图像张量与分批后的标签张量构成。默认的 collate_fn 有以下几种特性：

会在张量的首位处为批次创建一个维度
总是将 NumPy 数组与 PyTorch 数值型数据转换为 PyTorch 的 Tensor 类型
保持数据结构不变，比如，如果每个样本都是字典类型的，它会返回一个字典，该字典包含相同的键集，但是将批量化的张量作为值（如果原值集无法张量化，则将列表作为字典的值）。列表、元组以及命名元组等等与此类似。

用户可以使用自定义的 collate_fn。比如：沿着除第一个维度以外的维度进行整理、填充具有不同长度的序列、添加自定义类型的支持等等。

4. 单进程与多进程数据加载

DataLoader 默认使用单进程。
Python 中的 GIL 使线程的并发成为伪并发。为防止加载数据时代码发生阻塞，PyTorch 提供了一种简单的方式实现数据的多进程加载，使用时仅仅需要为参数 num_workers 设置一个正整数。

4.1. 单进程数据加载

在这种模式中，数据获取是在初始化 DataLoader 的同一个进程中完成的。因此，数据加载可能会阻塞计算。不过，在进程（共享内存、数据描述符等）间数据共享的场景中，或者在数据量小到足以完全加载进内存的场景中，这种模式可能会是比较合适的。此外，单进程的加载方式遇到异常中断时，提供的错误信息比较容易理解，方便调试。

4.2. 多进程数据加载

用正整数设置参数 num_workers 时，会返回一个有指定加载进程数量的多进程加载工具。
本模式中，每次都会创建一个 DataLoader 的迭代器（比如在调用 enumerate(dataloader) 时），也会创建 num_workers 个工作进程。dataset，collate_fn 以及 worker_init_fn 被传进每个进程中，进而被用于初始化和获取数据。这意味着连同内部IO在内的数据访问以及变换都在工作进程中执行。
在工作的子进程中执行 torch.utils.data.get_worker_info() 可以得到其信息（如进程ID，复制进去的数据集，初始化种子，等等），在 main 进程中调用的时候会返回 None。用户可以在数据集代码或者 worker_init_fn 中调用该函数，以保证每个子进程能独立的操作其自身的数据集实例副本，也可以用来确定某些代码是否运行在子进程中。比如，在对数据集分片时该函数会很有用。
对 Map 类型的数据集来说，主进程使用采样器生成序号列表，并把它们传进子进程。所以打乱顺序的操作是在主进程中完成的，之后主进程会通过分配加载的序号来引导加载过程。
对迭代器类型的数据集来说，因为每一个工作进程都会得到一个数据集实例的副本，这会导致简单的多进程产生重复的结果。通过使用 torch.utils.data.get_worker_info() 或 worker_init_fn，用户可以独立地配置每一个副本。（详见 IterableDataset 文档。）同理，在多进程加载中，参数 drop_last 会将每个进程中最后那个不满批次容量的批次舍弃掉。
迭代完成或者迭代器被垃圾回收之后，子进程就会被关闭。

Warning
不建议在多进程加载过程中返回 CUDA 张量，这是因为多进程中使用 CUDA 或者共享 CUDA 张量时会有很多精细的地方（参考多进程中的 CUDA）。多进程加载时，建议使用自动内存固定（即 pin_memory=True），它会将数据转到能使用 CUDA 的显卡上。

4.2.1. 平台上的不同表现

即便子进程依赖于 Python 的多进程，但在 Windows 与 Unix 平台上的表现仍旧是不同的。

在 Unix 上默认的多进程启动程序是 fork()。它使子进程可以直接通过复制的地址空间访问数据集和 Python 参数函数。
在 Windows 上默认的多进程启动程序是 spawn()，它会启动另一个解释器，并在其中执行主脚本，解释器之后是通过 pickle 模块序列化的方式接收 dataset、collate_fn以及其他参数的内部子进程函数。

这种单独序列化的方式意味着，为保证多进程数据加载时与Windows兼容，你应该采取两个步骤：

将主脚本用 if __name__ == '__main__': 包裹起来，以保证它在子进程被装载之后不会被多次执行。比如，你可以把数据集与 DataLoader 实例的创建逻辑放在这里。
确保自定义的 collate_fn，worker_init_fn 或者 dataset 的代码是在顶层的定义中并且是在 __main__ 校验的外面声明的，这能保证它们在子进程中的可用性。

4.2.2. 多进程数据加载中的随机性

默认的情况下，每个子进程使用 base_seed + worker_id 的方式设置随机种子，其中 base_seed 是由主进程使用 RNG 模块生成的长整型。不过在初始化子进程的过程中，其他库（比如 NumPy）的种子也会被复制进去，这会导致每个子进程返回的随机数是相同的。（解决方式参考这里）
在 worker_init_fn 中，你可以通过 torch.utils.data.get_worker_info().seed 或 torch.initial_seed() 获取 PyTorch 种子，并在数据加载前用于其他包的种子上。

5. 内存固定（pinning）

使用固定内存区（比如页锁定区）从主机往显卡复制数据时速度会更快，使用方式可参考这里。
在数据加载中，为 DataLoader 设置 pin_memory=True 会自动把选取的数据张量置于固定内存中，这样在向支持 CUDA 的显卡中转移数据时会更快。
默认的固定内存的逻辑仅能识别张量或者包含张量的 Map/Iterable 类型。如果批是一个自定义类型（collate_fn返回的是自定义类型的）或者其内部的每一项是自定义类型的，那固定化的逻辑无法识别它们，并且返回的结果也不是在内存中固定之后的。如果想实现对自定义类型批或数据类型的支持，需要在自定义类型中实现 pin_memory() 方法。
举个例子

# 自定义类
class SimpleCustomBatch:
    def __init__(self, data):
        transposed_data = list(zip(*data))
        self.inp = torch.stack(transposed_data[0], 0)
        self.tgt = torch.stack(transposed_data[1], 0)

    # 自定义内存 pinning 方法
    def pin_memory(self):
        self.inp = self.inp.pin_memory()
        self.tgt = self.tgt.pin_memory()
        return self

def collate_wrapper(batch):
    return SimpleCustomBatch(batch)

inps = torch.arange(10 * 5, dtype=torch.float32).view(10, 5)
tgts = torch.arange(10 * 5, dtype=torch.float32).view(10, 5)
dataset = TensorDataset(inps, tgts)

loader = DataLoader(dataset, batch_size=2, collate_fn=collate_wrapper,
                    pin_memory=True)

for batch_ndx, sample in enumerate(loader):
    print(sample.inp.is_pinned())
    print(sample.tgt.is_pinned())

官网的后半部分列举并介绍了数据集相关的类，其中为一部分类举出了使用案例。
具有使用案例的类有：

torch.utils.data.IterableDataset
torch.utils.data.BufferedShuffleDataset(dataset, buffer_size)
torch.utils.data.WeightedRandomSampler(weights, num_samples, replacement=True, generator=None)
torch.utils.data.BatchSampler(sampler, batch_size, drop_last)
torch.utils.data.distributed.DistributedSampler(dataset, num_replicas=None, rank=None, shuffle=True, seed=0, drop_last=False)