TORCH.UTILS.DATA 官方文档翻译解释

原文链接：

torch.utils.data — PyTorch 2.0 documentation

PyTorch data loading utility 的核心是 torch.utils.data.DataLoader 这个类，它把dataset整合为一个python的可迭代对象，它支持：

•  map-style and iterable-style datasets,

• customizing data loading order,

• automatic batching,

• single- and multi-process data loading,

• automatic memory 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)

函数说明

·dataset：加载的数据集。
·batch_size：批大小。
·shuffle：是否将数据打乱。
·sampler：样本抽样。
·batch_sampler: 批量抽样
·num_workers：使用多进程加载的进程数，0代表不使用多进程。
·collate_fn：如何将多个样本数据拼接成一个batch，一般使用默认的拼接方式即可。
·pin_memory：是否将数据保存在pin memory区，pin memory中的数据转到GPU会快一些。
·drop_last：dataset中的数据个数可能不是batch_size的整数倍，drop_last为True会将多出来不足一个batch的数据丢弃

原文如下：

Parameters(参数意思解释):
dataset (Dataset) – dataset from which to load the data.

batch_size (int, optional) – how many samples per batch to load (default: 1).

shuffle (bool, optional) – set to True to have the data reshuffled at every epoch (default: False).

sampler (Sampler or Iterable, optional) – defines the strategy to draw samples from the dataset. Can be any Iterable with __len__ implemented. If specified, shuffle must not be specified.

batch_sampler (Sampler or Iterable, optional) – like sampler, but returns a batch of indices at a time. Mutually exclusive with batch_size, shuffle, sampler, and drop_last.

num_workers (int, optional) – how many subprocesses to use for data loading. 0 means that the data will be loaded in the main process. (default: 0)

collate_fn (Callable, optional) – merges a list of samples to form a mini-batch of Tensor(s). Used when using batched loading from a map-style dataset.

pin_memory (bool, optional) – If True, the data loader will copy Tensors into device/CUDA pinned memory before returning them. If your data elements are a custom type, or your collate_fn returns a batch that is a custom type, see the example below.

drop_last (bool, optional) – set to True to drop the last incomplete batch, if the dataset size is not divisible by the batch size. If False and the size of dataset is not divisible by the batch size, then the last batch will be smaller. (default: False)

timeout (numeric, optional) – if positive, the timeout value for collecting a batch from workers. Should always be non-negative. (default: 0)

worker_init_fn (Callable, optional) – If not None, this will be called on each worker subprocess with the worker id (an int in [0, num_workers - 1]) as input, after seeding and before data loading. (default: None)

generator (torch.Generator, optional) – If not None, this RNG will be used by RandomSampler to generate random indexes and multiprocessing to generate base_seed for workers. (default: None)

prefetch_factor (int, optional, keyword-only arg) – Number of batches loaded in advance by each worker. 2 means there will be a total of 2 * num_workers batches prefetched across all workers. (default value depends on the set value for num_workers. If value of num_workers=0 default is None. Otherwise if value of num_workers>0 default is 2).

persistent_workers (bool, optional) – If True, the data loader will not shutdown the worker processes after a dataset has been consumed once. This allows to maintain the workers Dataset instances alive. (default: False)

pin_memory_device (str, optional) – the data loader will copy Tensors into device pinned memory before returning them if pin_memory is set to true.

A. DatasetTypes 本方法支持的数据集类型

dataset 是整个语句中最重要的参数，它决定了加载数据的具体对象。pytorch支持两种类型的数据：

• map-style datasets,

• iterable-style datasets.

以下是分别对此两种数据的介绍：

1. map-style datasets:

映射式数据集，特点是实现了__getitem__()和__len__()两种协议，将所有的数据存入内存，通过index映射的方式获取数据。优点是速度快；但是缺点是占用内存，大的数据集是无法使用。(Pytorch 自定义Dataset 的两种构建方式：Map VS Iterable - 知乎 (zhihu.com))。原官方文档举了一个 dataset[idx] 的例子。

2. Iterable-style datasets：

可迭代式数据集，是IterableDataset的subclass，实施的是__iter__()协议。这种数据集无需将整个数据集读取到内存中，通过覆盖写iter迭代器的方式实现流的形式输入数据。无需满足内存大于整个数据集，也无需知道全部数据的大小。官方文档此处给的例子是 iter(dataset)

注意：

将IterableDataset用于多线程任务的时候要避免重复取用数据，原文推荐详细阅读以下链接：torch.utils.data — PyTorch 2.0 documentation (torch.utils.data — PyTorch 2.0 documentationhttps://pytorch.org/docs/stable/data.html#torch.utils.data.IterableDatasettorch.utils.data — PyTorch 2.0 documentation (

B. Data Loading Order and Sampler数据加载的顺序和抽样

1. iterable-style datasets的数据加载顺序完全由user-defined iterable控制，它允许简单的chunk-reading(大文件读取)和dynamic batch size(每次输出一个批次的样本)

2. map-style dataset则是用 torch.utils.data.Sampler这个类来确定数据加载时使用的具体索引或键。比如在使用”随机梯度下山法“(SGD)的时候，Sampler这个类会随机排列出一个索引的列表，然后一次输出一个。或者在mini-batch SGD的情况时，一次输出一小批。

shuffle这个参数，可以决定是否对输出的数据做打乱操作。或者用sampler参数规定每次输出下一个索引。

注意：

sampler和batch_sampler都不适用于iterable-style datasets，因为在这种数据的情况下，没有index/key的概念。

C. Loading Batched and Non-Batched Data批量和非批量加载数据

DataLoader可以自动将收集的单个数据变为批量数据。使用的参数为：batch_size, drop_last, batch_sampler 和collate_fn 

1. Automatic batching(default) 自动批处理

自动批处理是最常见的情况，一般来说，张量的第一个维度就是批处理维度(batch dimension)

当batch_size(默认值为1)的值不为1的时候，DataLoader就会执行批处理。batch_size和drop_last这两个参数就会被用于确定data loader如何获得批量的dataset keys。对于map-style datasets, 还可以用batch_sampler(批量取样)参数来替代，这种情况下会一次生成一个keys的列表。

注意1：

batch_size和drop_last参数主要用于从sampler中构造一个batch_sampler。(1) 对于map-style datasets，sampler(采样器)由用户提供或者根据shuffle参数构造。(2) 对于Iterable-style datasets，sampler是一个虚拟的无限采样器。

注意2：

当使用多进程(multi-processing)从可迭代式数据集(iterable-style datasets)获取数据时，drop_last参数会丢弃每个工作进程数据集副本的最后一批非完整数据。

在使用sampler中的索引获取样本列表后，collate_fn 参数传递的函数将会把样本列表整理成batches。

在这种情况下，从map-style dataset 加载大致相当于下面的代码：

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

而从 iterable-style dataset加载数据，大致相当于下面的代码：

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

collate_fn可以自定义。

2. Disable automatic batching

在某些情况下，用户可能希望在数据集代码中手动处理批处理，或者直接加载单个样本。例如，直接加载批量数据（如从数据库批量读取数据或读取连续的内存块）可能cheaper，或者批量大小取决于数据，或者程序被设计为在单个样本上工作。在这些情况下，最好不要使用自动批处理（其中collate_fn用于整理样本），而是让数据加载器直接返回数据集对象的每个member。

把batch_size和batch_sampler都设置为None(batch_sampler的默认值已经为None，但batch_size这个参数要额外处理), 自动批处理会被禁用。从数据集中获取的每个样本都将使用作为collate_fn参数传递的函数进行处理。

当自动批处理被禁用时，默认的 collate_fn 简单地将 NumPy 数组转换为 PyTorch Tensors，而保持其他一切不变。

在这种情况下，从map-style dataset 加载大致相当于下面的代码：

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

而从 iterable-style dataset加载数据，大致相当于下面的代码：

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

3. Working with collate_fn:

启用或禁用自动批处理时，collate_fn的使用略有不同。

当自动批处理被禁用时，每个单独的数据样本都会调用collate_fn，并从数据加载器迭代器中输出。在这种情况下，默认的 collate_fn 简单地将 NumPy 数组转换为 PyTorch 张量。

当启用自动批处理时，每次调用collate_fn时都会包含一个数据样本列表。它会把输入样本整理成一个batch，然后从data loader iterater中输出。<本节的其余部分描述了默认 collate_fn (default_collate()) 的行为>

例如，如果每个数据样本由3个通道的图像和一个integral class label组成，即: 数据集的每个元素返回一个元组（image, class_index），那么默认的collate_fn会将这样的元组列表整理成一个由批次图像张量和批次类标签张量组成的单一元组(这段话说白了就是: 原本三通道彩色图像是一个三维矩阵，经过collate_fn处理以后，多个图像堆叠出了第四个维度)。

缺省的collate_fn具有以下特性：

1. 它总是预先添加一个新的维度作为批处理维度。

2. 自动将NumPy数组和Python数值转换为PyTorch张量。

3. 它保留数据结构，例如，如果每个样本是一个字典，它将输出一个具有相同键集的字典，但将批处理的Tensors作为值（如果值不能转换为Tensors，则输出列表）。对于list、tuple、namedtuple 等也是如此。

用户可以使用自定义的collate_fn来实现自定义批处理，例如，沿着第一个维度以外的维度进行批处理，填充不同长度的序列，或者添加对自定义数据类型的支持。

如果您遇到DataLoader输出的尺寸或类型与您的期望不同的情况，您可能需要检查您的collate_fn。

D. Single- and Multi-process Data Loading单进程和多进程数据加载

DataLoader默认使用单进程加载数据。

PyTorch 提供了一个简单的开关来执行多进程数据加载，只需将参数 num_workers 设置为一个正整数即可。

*注意一下，调试的时候合适使用单进程

*此处有一个关于多进程的警告，翻译如下：

-------------------------------------------------------------------------------------------------------------------

经过多次迭代后，加载器工作进程将消耗与父进程相同数量的 CPU 内存，用于从工作进程访问父进程中的所有 Python 对象。

如果Dataset包含大量数据（例如，在构建Dataset时加载了一个非常大的文件名列表）和/或使用了大量的工作进程（总体内存使用量为工作进程的数量*父进程的大小），这可能会产生问题。

最简单的解决方法是将Python对象替换为非反射计算的表示(non-refcounted representations)，例如Pandas、Numpy或PyArrow对象。查看问题 issue #13246 ，了解更多关于为什么会出现这种情况的细节，以及如何解决这些问题的示例代码。

--------------------------------------------------------------------------------------------------------------------

在此模式下，每次创建 DataLoader 的迭代器（例如，调用 enumerate(dataloader) 时），都会创建 num_workers 工作进程。此时，数据集、collate_fn和worker_init_fn将被传递给每个worker，它们将被用于初始化和获取数据。这意味着数据集访问及其内部 IO、转换（包括 collate_fn）在 Worker 进程中运行。

torrent.utils.data.get_worker_info()在worker进程中返回各种有用信息（包括worker id、数据集副本、初始种子等），在主进程中返回None。用户可以在数据集代码和/或worker_init_fn中使用该函数来单独配置每个数据集副本，并确定代码是否在worker进程中运行。（这在分割数据集时特别有用。）

对于map-style datasets，主进程使用采样器生成索引，并将其发送给worker。因此，任何洗牌随机化(shuffle)都是在主进程中完成的，主进程通过分配要加载的索引来指导加载。

对于可迭代式数据集，由于每个工作进程都会获得数据集对象的副本，因此简单实施多进程往往会导致数据重复。使用 torch.utils.data.get_worker_info() 和/或 worker_init_fn，用户可以独立配置(configure)每个副本。(See IterableDataset documentations for how to achieve this. ) 

出于类似的原因，在多进程加载中，drop_last 参数会丢弃每个 Worker 的可迭代式数据集副本的最后一批非完整数据。

一旦迭代结束，或者迭代器被垃圾回收，工作器就会被关闭。

警告：

--------------------------------------------------------------------------------------------------------------------

通常不建议在多进程加载中返回 CUDA 张量，因为在多进程中使用 CUDA 和共享 CUDA 张量存在许多微妙之处(see CUDA in multiprocessing)。相反，我们建议使用 automatic memory pinning (即设置pin_memory=True)，这样可以将数据快速传输到支持CUDA的GPU上。

--------------------------------------------------------------------------------------------------------------------

(a) Platform-specific behaviors 特定于平台的行为
由于 workers 依赖于  Python multiprocessing，因此 Windows 和Unix在启动worker上是不一样的。

1) 在 Unix 上，fork() 是默认的多进程启动方法。使用 fork()，子 worker 通常可以通过克隆的地址空间直接访问数据集和 Python 参数函数。

2) 在 Windows 或 MacOS 上，spawn() 是默认的多进程启动方法。使用 spawn()，另一个解释器被启动，运行你的主脚本，然后是内部 worker 函数，它通过 pickle 序列化(pickle serialization)接收数据集、collate_fn 和其它参数。

这种单独的序列化意味着您应该采取两个步骤来确保您在使用多进程数据加载时与Windows兼容：

1) 在 if __name__ == '__main__':块中封装主脚本的大部分代码，以确保在每个工作进程启动时不会再次运行（很可能产生错误）。您可以将数据集和 DataLoader 实例创建逻辑放在此处，因为它不需要在 worker 中重新执行。

2）确保任何自定义的 collate_fn、worker_init_fn 或数据集代码都声明为顶层定义( is declared as top level definitions)——放在对__main__ 的check之外。这将确保它们在 worker 进程中可用。(由于函数仅作为引用而不是字节码(bytecode)被拾取(pickled)，因此需要这样做）。

(b) Randomness in multi-process data loading

默认情况下，每个 Worker 的 PyTorch 种子设置为 base_seed + worker_id，其中 base_seed 是由主进程使用其 RNG 生成的 long（因此，强制消耗 RNG 状态）或指定的生成器。然而，其他库的种子可能会在初始化 Worker 时被复制，导致每个 Worker 返回相同的随机数。(参见FAQ中的这一部分）。

在 worker_init_fn 中，你可以使用 torch.utils.data.get_worker_info().seed 或 torch.initial_seed() 访问每个 worker 的 PyTorch 种子集，并在数据加载前使用它为其他库提供种子。

E. Memory Pinning

主机到GPU的拷贝如果来自pinned(page-locked)的内存，速度会更快。有关何时以及如何使用pinned内存的详细信息, See Use pinned memory buffers

对于数据加载，向DataLoader传递pin_memory=True将自动把获取的数据Tensors放在pinned内存中，从而能够更快地将数据传输到支持CUDA的GPU。

默认的内存钉住逻辑( pinning logic)只识别Tensors以及包含Tensors的映射和迭代表。默认情况下，如果钉住逻辑( pinning logic)看到一个自定义类型的批处理（如果您有一个返回自定义批处理类型的collate_fn，就会出现这种情况），或者如果您的批处理中的每个元素都是自定义类型，那么钉住逻辑将无法识别它们，并将返回该批处理（或这些元素），而不会钉住内存。要启用自定义批处理或数据类型的内存钉住功能，请在自定义类型上定义pin_memory()方法。

后面的example直接看原网页即可：torch.utils.data — PyTorch 2.0 documentation