pytorch集锦(7)-处理数据DataLoader和Dataset(4)

使用collate_fn

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

禁用自动批处理时，将对每个单独的数据样本调用collate_fn，并从数据加载器迭代器生成输出。在本例中，默认的collate_fn只是将NumPy数组转换为PyTorch张量。

启用自动批处理后，每次都会使用数据样本列表调用collate_fn。期望将输入样本整理成一批，以便从数据加载器迭代器中生成。本节的其余部分将描述默认的collate_fn（default_collate（））的行为。

例如，如果每个数据样本由一个3通道图像和一个完整的类标签组成，即数据集的每个元素都返回一个元组（image，class_index），则默认的collated_fn将这些元组的列表整理成一个由批处理图像张量和批处理类标签tensor组成的单个元组。特别是，默认的collate_fn具有以下属性：

它总是在前面添加一个新维度作为批处理维度。

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

它保留了数据结构，例如，如果每个样本都是一个字典，它输出一个具有相同密钥集的字典，但将张量作为值进行批处理（如果值不能转换为张量，则列出）。列表、元组、命名元组等也是如此。

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

如果遇到DataLoader的输出具有不同于预期的维度或类型的情况，您可能需要检查您的collate_fn。

单进程和多进程数据加载

默认情况下，DataLoader使用单进程数据加载。

在Python进程中，全局解释器锁（GIL）阻止跨线程真正完全并行化Python代码。为了避免在数据加载时阻塞计算代码，PyTorch提供了一个简单的开关来执行多进程数据加载，只需将参数num_workers设置为正整数。

单进程数据加载（默认）

在此模式下，数据获取在初始化DataLoader的相同过程中完成。因此，数据加载可能会阻碍计算。然而，当用于在进程之间共享数据的资源（例如，共享内存、文件描述符）受到限制时，或者当整个数据集很小并且可以完全加载到内存中时，这种模式可能是优选的。此外，单进程加载通常显示更可读的错误跟踪，因此对调试非常有用。

多进程数据加载

将参数num_workers设置为正整数将启用具有指定数量的加载器工作进程的多进程数据加载。

在多次迭代之后，加载器工作进程将为父进程中从工作进程访问的所有Python对象消耗与父进程相同数量的CPU内存。如果数据集包含大量数据（例如，在数据集构建时加载了一个非常大的文件名列表）和/或使用了大量工作者（总内存使用量是工作人员的数量*父进程的大小），这可能会有问题。最简单的解决方法是用非引用表示（如Pandas、Numpy或PyArrow对象）替换Python对象。

了解更多关于为什么会出现这种情况的详细信息，以及如何解决这些问题的示例代码：https://github.com/pytorch/pytorch/issues/13246

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

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

对于映射样式的数据集，主要过程使用采样器生成索引并将其发送给工作人员。因此，任何随机排序都是在主过程中完成的，该过程通过为加载分配索引来引导加载。

对于可迭代样式的数据集，由于每个工作进程都会获得数据集对象的副本，因此天真的多进程加载通常会导致重复数据。使用torch.utils.data.get_worker_info（）和/或worker_init_fn，用户可以独立配置每个副本。（有关如何实现这一点，请参阅IterableDataset文档。）出于类似的原因，在多进程加载中，drop_last参数会删除每个工作程序的可迭代样式数据集副本的最后一个非完整批。

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

通常不建议在多进程加载中返回CUDA张量，因为在多处理中使用CUDA和共享CUDA张量有许多微妙之处（请参阅多处理中的CUDA）。相反，我们建议使用自动内存固定（即，将pin_memory=True设置为True），这可以实现到启用CUDA的GPU的快速数据传输。

平台特定行为

由于工作者依赖于Python多处理，因此与Unix相比，Windows上的工作人员启动行为有所不同。

在Unix上，fork（）是默认的多处理启动方法。使用fork（），子工作人员通常可以通过克隆的地址空间直接访问数据集和Python参数函数。

在Windows或MacOS上，spawn（）是默认的多处理启动方法。使用spawn（），将启动另一个解释器，该解释器运行主脚本，然后是内部worker函数，该函数通过pickle序列化接收数据集、整理fn和其他参数。

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

将主脚本的大部分代码包装在if name==‘main’：块中，以确保在启动每个工作进程时不会再次运行（很可能会产生错误）。您可以将数据集和DataLoader实例创建逻辑放在这里，因为它不需要在worker中重新执行。

确保在__main__检查之外，将任何自定义的collated_fn、worker_init_fn或数据集代码声明为顶级定义。这确保了它们在工作进程中可用。（这是必需的，因为函数仅作为引用而不是字节码。）

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

默认情况下，每个工作进程将其PyTorch种子设置为base_sed+worker_id，其中base_sed是由主进程使用其RNG（因此，强制使用RNG状态）或指定的生成器生成的长消息。然而，其他库的种子可能会在初始化工作者时复制，导致每个工作者返回相同的随机数。（请参见常见问题解答中的本节）。

在worker_init_fn中，可以使用torch.utils.data.get_worker_info（）.seed或torch.initial_sed（）访问每个worker的PyTorch种子集，并在数据加载之前使用它为其他库种子。

内存固定(Memory Pinning)

当主机到GPU的复制源于固定（页面锁定）内存时，速度要快得多。有关何时以及如何通常使用固定内存的详细信息，请参见使用固定内存缓冲区。

对于数据加载，将pin_memory=True传递给DataLoader将自动将获取的数据Tensor放入固定内存中，从而实现更快的数据传输到支持CUDA的GPU。

默认的内存锁定逻辑只识别张量和包含张量的映射和可迭代项。默认情况下，如果固定(pinning)逻辑看到一个自定义类型的批（如果您有一个返回自定义批类型的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)

    # custom memory pinning method on custom type
    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.DataLoader

class
torch.utils.data.DataLoader(dataset, batch_size=1, shuffle=None, 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, generator=None, *, prefetch_factor=2, persistent_workers=False, pin_memory_device='')

数据加载器。组合数据集和采样器，并在给定数据集上提供可迭代的。

DataLoader支持具有单进程或多进程加载、自定义加载顺序以及可选的自动批处理（排序）和内存固定的映射样式和可迭代样式数据集。

参数：

dataset (Dataset) –从中加载数据的数据集。

batch_size (int, optional) –每个批次要加载多少个样本（默认值：1）。

shuffle (bool, optional) –设置为True以在每个历元重新排列数据（默认值：False）。

sampler (Sampler or Iterable, optional)–定义从数据集中抽取样本的策略。可以是实现了__len__的任何Iterable。如果已指定，则不得指定shuffle。

batch_sampler (Sampler or Iterable, optional) ––类似采样器，但一次返回一批索引。与batch_size、shuffle、sampler和drop_last互斥。

num_workers (int, optional) –用于数据加载的子流程数。0表示将在主进程中加载数据。（默认值：0）

collate_fn (Callable, optional)–合并样本列表以形成一个小批量Tensor。在使用从地图样式数据集批量加载时使用。

pin_memory (bool, optional) –如果为True，数据加载器将在返回Tensor之前将其复制到设备/CUDA固定内存中。如果您的数据元素是自定义类型，或者您的collate_fn返回一个自定义类型的批处理，请参见下面的示例。

drop_last (bool, optional) –如果数据集大小不能被批大小整除，则设置为True以删除最后一个不完整的批。如果False，并且数据集的大小不能被批大小整除，则最后一批将更小。（默认值：False）

timeout (numeric, optional) –如果为正值，则为从工作人员收集批次的超时值。应始终为非负。（默认值：0）

worker_init_fn（可调用，可选）–如果不是None，则将在种子设定之后和数据加载之前，以worker id（[0，num_workers-1]中的一个int）作为输入，在每个worker子进程上调用此函数。（默认值：无）

generator（torch.generator，可选）–如果不是None，RandomSampler将使用此RNG生成随机索引，并进行多处理以生成工作线程的base_sed。（默认值：无）

prefetch_factor（int，可选，仅限关键字arg）–每个工作程序预先加载的批处理数。2表示在所有工作进程中总共预取了2*num_workers批。（默认值：2）

persistent_workers（bool，可选）–如果为True，则数据加载器在数据集消耗一次后不会关闭工作进程。这允许保持工人数据集实例处于活动状态。（默认值：False）

pin_memory_device（str，可选）–如果pin_memory设置为true，数据加载器将在返回Tensor之前将Tensor复制到设备固定内存中。

如果使用派生启动方法，worker_init_fn不能是不可拾取的对象，例如lambda函数。

len(dataloader)启发式基于所使用的采样器的长度。当数据集是IterableDataset时，它会返回基于len（dataset）/batch_size的估计值，并根据drop_last进行适当舍入，而不考虑多进程加载配置。这是PyTorch能够做出的最好猜测，因为PyTorch信任用户数据集代码正确处理多进程加载以避免重复数据。

然而，如果分片导致多个工人的最后一批不完整，则此估计仍可能不准确，因为（1）否则完整的批次可能会被拆分为多个批次，（2）当设置drop_last时，可能会丢弃一批以上的样本。不幸的是，PyTorch通常无法检测到此类病例。

torch.utils.data.Dataset

class
torch.utils.data.Dataset(*args, **kwds)

表示数据集的抽象类。

表示从键到数据样本的映射的所有数据集都应该将其子类化。所有子类都应该覆盖__getitem_（），支持获取给定键的数据样本。子类还可以选择性地覆盖__len__（），许多Sampler实现和DataLoader的默认选项都期望它返回数据集的大小。

默认情况下，DataLoader构造一个生成整数索引的索引采样器。要使其与具有非整数索引/键的地图样式数据集一起工作，必须提供自定义采样器。

torch.utils.data.IterableDataset

class torch.utils.data.IterableDataset(*args, **kwds)

可迭代数据集。

所有表示数据样本可迭代的数据集都应该将其子类化。当数据来自流时，这种形式的数据集特别有用。

所有子类都应该覆盖__iter_（），这将返回此数据集中样本的迭代器。

当子类与DataLoader一起使用时，数据集中的每个项都将从DataLoader迭代器生成。当num_workers>0时，每个工作进程将具有数据集对象的不同副本，因此通常需要单独配置每个副本，以避免从工作进程返回重复数据。在工作进程中调用get_worker_info（）时，返回有关工作进程的信息。它可以在数据集的__iter_（）方法或DataLoader的worker_init_fn选项中使用，以修改每个副本的行为。