torch.utils.data#

创建于: 2025年6月13日 | 最后更新于: 2025年6月13日

PyTorch 数据加载工具的核心是 torch.utils.data.DataLoader 类。它代表了一个可迭代的数据集，支持

map-style 和 iterable-style 数据集,
自定义数据加载顺序,
自动批处理,
单进程和多进程数据加载,
自动内存锁定.

这些选项由 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)

下面的章节将详细介绍这些选项的效果和用法。

数据集类型#

DataLoader 构造函数最重要的参数是 dataset，它表示要从中加载数据的数据集对象。PyTorch 支持两种不同的数据集类型：

Map-style 数据集,
Iterable-style 数据集.

Map-style 数据集#

Map-style 数据集是指实现了 __getitem__() 和 __len__() 协议的数据集，它将（可能非整数的）索引/键映射到数据样本。

例如，这样的数据集在通过 dataset[idx] 访问时，可以从磁盘上的文件夹读取第 idx 张图像及其对应的标签。

有关更多详细信息，请参阅 Dataset。

Iterable-style 数据集#

Iterable-style 数据集是 IterableDataset 的子类实例，它实现了 __iter__() 协议，并表示数据样本的可迭代序列。这种类型的数据集特别适合随机读取昂贵或几乎不可能，且批次大小取决于所获取数据的情况。

例如，当调用 iter(dataset) 时，这种数据集可以返回一个从数据库、远程服务器甚至实时生成的日志中读取的数据流。

有关更多详细信息，请参阅 IterableDataset。

注意

当使用 IterableDataset 进行多进程数据加载时，同一个数据集对象会在每个工作进程中复制，因此必须对副本进行不同的配置以避免重复数据。有关如何实现这一点，请参阅 IterableDataset 文档。

数据加载顺序与 `Sampler`#

对于 iterable-style 数据集，数据加载顺序完全由用户定义的迭代器控制。这允许更轻松地实现分块读取和动态批次大小（例如，每次生成一个批次的样本）。

本节的其余部分涉及 map-style 数据集的情况。torch.utils.data.Sampler 类用于指定数据加载中使用的索引/键的序列。它们表示数据集索引的可迭代对象。例如，在随机梯度下降（SGD）的常见情况下，Sampler 可以随机排列索引列表并一次生成一个索引，或者为 mini-batch SGD 生成少量索引。

顺序或打乱的 sampler 将根据 DataLoader 的 shuffle 参数自动构造。或者，用户可以使用 sampler 参数指定一个自定义的 Sampler 对象，该对象每次生成要获取的下一个索引/键。

一次生成一批索引的自定义 Sampler 可以作为 batch_sampler 参数传递。也可以通过 batch_size 和 drop_last 参数启用自动批处理。有关更多详细信息，请参阅下一节。

注意

Neither sampler nor batch_sampler is compatible with iterable-style datasets, since such datasets have no notion of a key or an index.

加载批处理和非批处理数据#

DataLoader 支持通过 batch_size、drop_last、batch_sampler 和 collate_fn（具有默认函数）参数自动将单独获取的数据样本组合成批次。

自动批处理（默认）#

这是最常见的情况，对应于获取数据的小批量（minibatch）并将其组合成批处理样本，即包含 Tensor，其中一个维度是批次维度（通常是第一个）。

当 batch_size（默认为 1）不为 None 时，数据加载器将生成批处理样本而不是单独的样本。batch_size 和 drop_last 参数用于指定数据加载器如何获取数据集键的批次。对于 map-style 数据集，用户还可以指定 batch_sampler，它一次生成一个键列表。

注意

The batch_size and drop_last arguments essentially are used to construct a batch_sampler from sampler. For map-style datasets, the sampler is either provided by user or constructed based on the shuffle argument. For iterable-style datasets, the sampler is a dummy infinite one. See this section on more details on samplers.

注意

When fetching from iterable-style datasets with multi-processing the drop_last argument drops the last non-full batch of each worker’s dataset replica.

After fetching a list of samples using the indices from sampler, the function passed as the collate_fn argument is used to collate lists of samples into batches.

In this case, loading from a map-style dataset is roughly equivalent with

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

and loading from an iterable-style dataset is roughly equivalent with

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

A custom collate_fn can be used to customize collation, e.g., padding sequential data to max length of a batch. See this section on more about collate_fn.

禁用自动批处理#

在某些情况下，用户可能希望在数据集代码中手动处理批处理，或者只是加载单独的样本。例如，直接加载批处理数据（例如，从数据库批量读取或读取连续内存块）可能更便宜，或者批次大小依赖于数据，或者程序设计为处理单个样本。在这些情况下，最好不要使用自动批处理（其中 collate_fn 用于组合样本），而是让数据加载器直接返回 dataset 对象的每个成员。

When both batch_size and batch_sampler are None (default value for batch_sampler is already None), automatic batching is disabled. Each sample obtained from the dataset is processed with the function passed as the collate_fn argument.

When automatic batching is disabled, the default collate_fn simply converts NumPy arrays into PyTorch Tensors, and keeps everything else untouched.

In this case, loading from a map-style dataset is roughly equivalent with

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

and loading from an iterable-style dataset is roughly equivalent with

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

See this section on more about collate_fn.

使用 `collate_fn`#

The use of collate_fn is slightly different when automatic batching is enabled or disabled.

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

启用自动批处理时，collate_fn 每次都会与一个数据样本列表一起调用。它需要将输入的样本进行合并，以构成一个批次，然后从数据加载器迭代器中产生。本节的其余部分将描述默认 collate_fn (default_collate()) 的行为。

例如，如果每个数据样本由一个 3 通道图像和一个整数类标签组成，也就是说，数据集的每个元素返回一个元组 (image, class_index)，则默认的 collate_fn 会将这样的一组元组合并成一个包含批量图像张量和批量类标签张量的单个元组。特别是，默认的 collate_fn 具有以下属性：

它始终添加一个新维度作为批次维度。
它会自动将 NumPy 数组和 Python 数值转换为 PyTorch 张量。
它会保留数据结构，例如，如果每个样本都是一个字典，它将输出一个具有相同键集的字典，但值为批量张量（如果值无法转换为张量，则为列表）。对于 list、tuple、namedtuple 等也是如此。

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

如果您发现 DataLoader 的输出的维度或类型与您的预期不同，您可能需要检查您的 collate_fn。

单进程和多进程数据加载#

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

在 Python 进程中，全局解释器锁 (GIL) 阻止了跨线程的真正完全并行的 Python 代码。为了避免数据加载阻塞计算代码，PyTorch 提供了一个简单的切换方式，通过将参数 num_workers 设置为正整数来执行多进程数据加载。

单进程数据加载（默认）#

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

多进程数据加载#

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

警告

经过几次迭代后，加载器工作进程将消耗与父进程相同的 CPU 内存，用于父进程中从工作进程访问的所有 Python 对象。如果 Dataset 包含大量数据（例如，您在 Dataset 构建时加载了非常大的文件名列表）和/或您使用了大量工作进程（总体内存使用量为 工作进程数 * 父进程大小），这可能会有问题。最简单的解决方法是将 Python 对象替换为非引用计数的表示形式，如 Pandas、Numpy 或 PyArrow 对象。有关此问题发生的原因以及如何解决这些问题的示例代码，请参阅 issue #13246。

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

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

对于 map-style Dataset，主进程使用 sampler 生成索引，并将它们发送给工作进程。因此，任何 shuffle 随机化都在主进程中完成，它通过分配要加载的索引来指导加载。

对于 iterable-style Dataset，由于每个工作进程都获得 Dataset 对象的一个副本，因此简单的多进程加载通常会导致重复数据。使用 torch.utils.data.get_worker_info() 和/或 worker_init_fn，用户可以独立配置每个副本。（有关如何实现此目的，请参阅 IterableDataset 文档。）出于类似的原因，在多进程加载中，drop_last 参数会丢弃每个工作进程的可迭代 Dataset 副本的最后一个不完整批次。

一旦达到迭代的末尾，或者迭代器被垃圾回收，工作进程就会被关闭。

警告

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

平台特定的行为#

由于工作进程依赖于 Python 的 multiprocessing，因此工作进程的启动行为在 Windows 上与 Unix 不同。

在 Unix 上，fork() 是默认的 multiprocessing 启动方法。使用 fork()，子工作进程通常可以通过克隆的地址空间直接访问 dataset 和 Python 参数函数。
在 Windows 或 MacOS 上，spawn() 是默认的 multiprocessing 启动方法。使用 spawn()，会启动另一个解释器来运行您的主脚本，然后是内部工作函数，该函数通过 pickle 序列化接收 dataset、collate_fn 和其他参数。

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

将您主脚本的大部分代码包装在 if __name__ == '__main__': 块中，以确保它在启动每个工作进程时不会再次运行（很可能会产生错误）。您可以在这里放置您的 Dataset 和 DataLoader 实例创建逻辑，因为它们不需要在工作进程中重新执行。
确保任何自定义的 collate_fn、worker_init_fn 或 dataset 代码都被定义为顶层定义，位于 __main__ 检查之外。这可以确保它们在工作进程中可用。（这是必需的，因为函数仅作为引用被 pickle，而不是作为 bytecode。）

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

默认情况下，每个工作进程的 PyTorch 种子将设置为 base_seed + worker_id，其中 base_seed 是由主进程使用其 RNG 生成的长整型（因此，强制消耗一个 RNG 状态）或指定的 generator。但是，其他库的种子在初始化工作进程时可能会重复，导致每个工作进程返回相同的随机数。（请参阅 FAQ 中的此节）。

在 worker_init_fn 中，您可以使用 torch.utils.data.get_worker_info().seed 或 torch.initial_seed() 来访问为每个工作进程设置的 PyTorch 种子，并使用它在数据加载之前为其他库设置种子。

内存固定#

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

对于数据加载，将 pin_memory=True 传递给 DataLoader 将自动将获取的数据张量复制到固定内存中，从而实现到 CUDA 启用 GPU 的更快数据传输。

默认的内存固定逻辑仅识别张量以及包含张量的映射和可迭代对象。默认情况下，如果固定逻辑遇到自定义类型的批次（当您使用返回自定义批次类型的 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())

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=None, persistent_workers=False, pin_memory_device='', in_order=True)[source]#

数据加载器结合了 Dataset 和 Sampler，并为给定的 Dataset 提供了一个可迭代对象。

DataLoader 支持 map-style 和 iterable-style Datasets，以及单进程或多进程加载、自定义加载顺序以及可选的自动批处理（collation）和内存固定。

有关更多详细信息，请参阅 torch.utils.data 文档页面。

参数

dataset (Dataset) – 用于加载数据的 Dataset。
batch_size (int, optional) – 每次加载多少个样本（默认为 1）。
shuffle (bool, optional) – 设置为 True 可在每个 epoch 重新排列数据（默认为 False）。
sampler (Sampler or Iterable, optional) – 定义从 Dataset 中抽取样本的策略。可以是任何实现了 __len__ 的 Iterable。如果指定了此参数，则必须不指定 shuffle。
batch_sampler (Sampler or Iterable, optional) – 类似于 sampler，但一次返回一个批次的索引。与 batch_size、shuffle、sampler 和 drop_last 互斥。
num_workers (int, optional) – 用于数据加载的子进程数量。 0 表示数据将在主进程中加载。（默认为 0）
collate_fn (Callable, optional) – 将样本列表合并以形成张量（或张量）的 mini-batch。在从 map-style Dataset 使用批处理加载时使用。
pin_memory (bool, optional) – 如果为 True，则数据加载器将在返回张量之前将它们复制到设备/CUDA 固定的内存中。如果您的数据元素是自定义类型，或者您的 collate_fn 返回的批次是自定义类型，请参见下面的示例。
drop_last (bool, optional) – 如果 Dataset 大小不能被 batch_size 整除，则设置为 True 以丢弃最后一个不完整的批次。如果为 False 且 Dataset 大小不能被 batch_size 整除，则最后一个批次将较小。（默认为 False）
timeout (numeric, optional) – 如果为正数，则表示从工作进程收集批次的超时值。应始终为非负数。（默认为 0）
worker_init_fn (Callable, optional) – 如果不为 None，则在每个工作进程子进程中调用此函数，并将工作进程 ID（[0, num_workers - 1] 中的整数）作为输入，在设置种子和数据加载之后。（默认为 None）
multiprocessing_context (str or multiprocessing.context.BaseContext, optional) – 如果为 None，则将使用您操作系统默认的 multiprocessing context。 (default: None)
generator (torch.Generator, optional) – 如果不为 None，则 RandomSampler 将使用此 RNG 来生成随机索引，并且 multiprocessing 将使用此 RNG 来为工作进程生成 base_seed。（默认为 None）
prefetch_factor (int, optional, keyword-only arg) – 每个工作进程预先加载的批次数。 2 表示所有工作进程总共将预加载 2 * num_workers 个批次。（默认值取决于 num_workers 的设置值。如果 num_workers=0，则默认为 None。否则，如果 num_workers > 0，则默认为 2）。
persistent_workers (bool, optional) – 如果为 True，则数据加载器不会在数据集被消耗一次后关闭工作进程。这允许保持工作进程的 Dataset 实例处于活动状态。（默认为 False）
pin_memory_device (str, optional) – 如果 pin_memory 为 True，则用于 pin_memory 的设备。如果未指定，则默认使用当前加速器。此参数不推荐使用，并可能被弃用。
in_order (bool, optional) – 如果为 False，则数据加载器不会强制要求批次以先入先出的顺序返回。仅当 num_workers > 0 时才适用。（默认为 True）

警告

如果使用 spawn 启动方法，则 worker_init_fn 不能是无法 pickle 的对象，例如 lambda 函数。有关 PyTorch 中多进程的最佳实践，请参阅多进程最佳实践。

警告

len(dataloader) 的启发式方法基于所使用的 sampler 的长度。当 Dataset 是 IterableDataset 时，它会根据 len(dataset) / batch_size 返回一个估计值，并根据 drop_last 进行适当的舍入，无论多进程加载配置如何。这是 PyTorch 能够做出的最佳猜测，因为 PyTorch 信任用户 dataset 代码正确处理多进程加载以避免重复数据。

但是，如果分片导致多个工作进程拥有不完整的最后一个批次，这个估计值仍然可能不准确，因为（1）一个本来完整的批次可能会被分成多个批次，以及（2）当设置 drop_last 时，可能会丢弃多个批次大小的样本。不幸的是，PyTorch 通常无法检测到这种情况。

有关这两种 Dataset 类型以及 IterableDataset 如何与多进程数据加载交互的更多信息，请参阅 Dataset 类型。

警告

有关与随机种子相关的问题，请参阅可复现性、我的数据加载器工作进程返回相同的随机数和多进程数据加载中的随机性。

警告

将 in_order 设置为 False 可能会影响可复现性，并且在数据不平衡的情况下可能导致将有偏斜的数据分布馈送到训练器。

class torch.utils.data.Dataset[source]#

表示 Dataset 的抽象类。

所有表示键到数据样本映射的 Dataset 都应该继承此类。所有子类都应重写 __getitem__()，支持根据给定键获取数据样本。子类也可以选择重写 __len__()，许多 Sampler 实现和 DataLoader 的默认选项期望此方法返回 Dataset 的大小。子类也可以选择实现 __getitems__()，以加速批次样本加载。此方法接受样本批次的索引列表，并返回样本列表。

注意

DataLoader 默认构建一个生成整数索引的索引采样器。为了使其与非整数索引/键的 map-style Dataset 一起工作，必须提供自定义采样器。

class torch.utils.data.IterableDataset[source]#

可迭代 Dataset。

所有表示数据样本可迭代对象的 Dataset 都应继承此类。这种形式的 Dataset 在数据来自流时尤其有用。

所有子类都应重写 __iter__()，它将返回此 Dataset 中样本的迭代器。

当子类与 DataLoader 一起使用时，Dataset 中的每个项都将从 DataLoader 迭代器中产生。当 num_workers > 0 时，每个工作进程将拥有 Dataset 对象的一个不同副本，因此通常希望独立配置每个副本以避免从工作进程返回重复数据。get_worker_info() 在工作进程中调用时，会返回有关工作进程的信息。它可以在 Dataset 的 __iter__() 方法或 DataLoader 的 worker_init_fn 选项中使用，以修改每个副本的行为。

示例 1：在 __iter__() 中拆分所有工作进程的负载

>>> class MyIterableDataset(torch.utils.data.IterableDataset):
...     def __init__(self, start, end):
...         super(MyIterableDataset).__init__()
...         assert end > start, "this example code only works with end >= start"
...         self.start = start
...         self.end = end
...
...     def __iter__(self):
...         worker_info = torch.utils.data.get_worker_info()
...         if worker_info is None:  # single-process data loading, return the full iterator
...             iter_start = self.start
...             iter_end = self.end
...         else:  # in a worker process
...             # split workload
...             per_worker = int(math.ceil((self.end - self.start) / float(worker_info.num_workers)))
...             worker_id = worker_info.id
...             iter_start = self.start + worker_id * per_worker
...             iter_end = min(iter_start + per_worker, self.end)
...         return iter(range(iter_start, iter_end))
...
>>> # should give same set of data as range(3, 7), i.e., [3, 4, 5, 6].
>>> ds = MyIterableDataset(start=3, end=7)

>>> # Single-process loading
>>> print(list(torch.utils.data.DataLoader(ds, num_workers=0)))
[tensor([3]), tensor([4]), tensor([5]), tensor([6])]

>>> # Multi-process loading with two worker processes
>>> # Worker 0 fetched [3, 4].  Worker 1 fetched [5, 6].
>>> print(list(torch.utils.data.DataLoader(ds, num_workers=2)))
[tensor([3]), tensor([5]), tensor([4]), tensor([6])]

>>> # With even more workers
>>> print(list(torch.utils.data.DataLoader(ds, num_workers=12)))
[tensor([3]), tensor([5]), tensor([4]), tensor([6])]

示例 2：使用 worker_init_fn 拆分所有工作进程的负载

>>> class MyIterableDataset(torch.utils.data.IterableDataset):
...     def __init__(self, start, end):
...         super(MyIterableDataset).__init__()
...         assert end > start, "this example code only works with end >= start"
...         self.start = start
...         self.end = end
...
...     def __iter__(self):
...         return iter(range(self.start, self.end))
...
>>> # should give same set of data as range(3, 7), i.e., [3, 4, 5, 6].
>>> ds = MyIterableDataset(start=3, end=7)

>>> # Single-process loading
>>> print(list(torch.utils.data.DataLoader(ds, num_workers=0)))
[3, 4, 5, 6]
>>>
>>> # Directly doing multi-process loading yields duplicate data
>>> print(list(torch.utils.data.DataLoader(ds, num_workers=2)))
[3, 3, 4, 4, 5, 5, 6, 6]

>>> # Define a `worker_init_fn` that configures each dataset copy differently
>>> def worker_init_fn(worker_id):
...     worker_info = torch.utils.data.get_worker_info()
...     dataset = worker_info.dataset  # the dataset copy in this worker process
...     overall_start = dataset.start
...     overall_end = dataset.end
...     # configure the dataset to only process the split workload
...     per_worker = int(math.ceil((overall_end - overall_start) / float(worker_info.num_workers)))
...     worker_id = worker_info.id
...     dataset.start = overall_start + worker_id * per_worker
...     dataset.end = min(dataset.start + per_worker, overall_end)
...

>>> # Mult-process loading with the custom `worker_init_fn`
>>> # Worker 0 fetched [3, 4].  Worker 1 fetched [5, 6].
>>> print(list(torch.utils.data.DataLoader(ds, num_workers=2, worker_init_fn=worker_init_fn)))
[3, 5, 4, 6]

>>> # With even more workers
>>> print(list(torch.utils.data.DataLoader(ds, num_workers=12, worker_init_fn=worker_init_fn)))
[3, 4, 5, 6]

class torch.utils.data.TensorDataset(*tensors)[source]#

包装张量的 Dataset。

每个样本将通过沿第一个维度索引张量来检索。

参数: *tensors (Tensor) – 具有相同第一个维度的张量。

class torch.utils.data.StackDataset(*args, **kwargs)[source]#

堆叠多个 Dataset 的 Dataset。

此类在将复杂的输入数据（以 Dataset 形式提供）的不同部分组合时很有用。

示例

>>> images = ImageDataset()
>>> texts = TextDataset()
>>> tuple_stack = StackDataset(images, texts)
>>> tuple_stack[0] == (images[0], texts[0])
>>> dict_stack = StackDataset(image=images, text=texts)
>>> dict_stack[0] == {'image': images[0], 'text': texts[0]}

参数

*args (Dataset) – 要堆叠的 Dataset，作为元组返回。
**kwargs (Dataset) – 要堆叠的 Dataset，作为字典返回。

class torch.utils.data.ConcatDataset(datasets)[source]#

拼接多个 Dataset 的 Dataset。

此类在组合不同的现有 Dataset 时很有用。

参数: datasets (sequence) – 要拼接的 Dataset 列表

class torch.utils.data.ChainDataset(datasets)[source]#

链接多个 IterableDataset 的 Dataset。

此类在组合不同的现有 Dataset 流时很有用。链接操作是即时完成的，因此使用此类连接大规模 Dataset 将是高效的。

参数: datasets (iterable of IterableDataset) – 要链接在一起的 Dataset

class torch.utils.data.Subset(dataset, indices)[source]#

指定索引处的 Dataset 的子集。

参数

dataset (Dataset) – 整个 Dataset
indices (sequence) – 在整个数据集中选取的子集索引

torch.utils.data._utils.collate.collate(batch, *, collate_fn_map=None)[source]#

通用 collate 函数，处理批次内收集类型的元素。

该函数还提供了函数注册表来处理特定的元素类型。default_collate_fn_map 为张量、numpy 数组、数字和字符串提供默认的 collate 函数。

参数

batch – 要进行 collate 的单个批次
collate_fn_map (Optional[dict[Union[type, tuple[type, ...]], Callable]]) – 可选的字典，将元素类型映射到相应的 collate 函数。如果元素类型不在此字典中，此函数将遍历字典中的每个键（按插入顺序），如果元素类型是键的子类，则调用相应的 collate 函数。

示例

>>> def collate_tensor_fn(batch, *, collate_fn_map):
...     # Extend this function to handle batch of tensors
...     return torch.stack(batch, 0)
>>> def custom_collate(batch):
...     collate_map = {torch.Tensor: collate_tensor_fn}
...     return collate(batch, collate_fn_map=collate_map)
>>> # Extend `default_collate` by in-place modifying `default_collate_fn_map`
>>> default_collate_fn_map.update({torch.Tensor: collate_tensor_fn})

注意

每个 collate 函数都需要一个用于批次的 positional 参数，以及一个用于 collate 函数字典的关键字参数，如 collate_fn_map。

torch.utils.data.default_collate(batch)[source]#

获取一个数据批次，并将批次中的元素放入一个张量中，并添加一个额外的维度——批次大小。

确切的输出类型可以是 torch.Tensor、Sequence of torch.Tensor、Collection of torch.Tensor，或者根据输入类型保持不变。当 DataLoader 中定义了 batch_size 或 batch_sampler 时，此函数用作 collate 的默认函数。

以下是输入类型（基于批次内元素的类型）到输出类型的通用映射：

torch.Tensor -> torch.Tensor （添加了批次大小的外部维度）

NumPy 数组 -> torch.Tensor

float -> torch.Tensor

int -> torch.Tensor

str -> str （不变）

bytes -> bytes （不变）

Mapping[K, V_i] -> Mapping[K, default_collate([V_1, V_2, …])]

NamedTuple[V1_i, V2_i, …] -> NamedTuple[default_collate([V1_1, V1_2, …]), default_collate([V2_1, V2_2, …]), …]

Sequence[V1_i, V2_i, …] -> Sequence[default_collate([V1_1, V1_2, …]), default_collate([V2_1, V2_2, …]), …]

参数: batch – 要进行 collate 的单个批次

示例

>>> # Example with a batch of `int`s:
>>> default_collate([0, 1, 2, 3])
tensor([0, 1, 2, 3])
>>> # Example with a batch of `str`s:
>>> default_collate(['a', 'b', 'c'])
['a', 'b', 'c']
>>> # Example with `Map` inside the batch:
>>> default_collate([{'A': 0, 'B': 1}, {'A': 100, 'B': 100}])
{'A': tensor([  0, 100]), 'B': tensor([  1, 100])}
>>> # Example with `NamedTuple` inside the batch:
>>> Point = namedtuple('Point', ['x', 'y'])
>>> default_collate([Point(0, 0), Point(1, 1)])
Point(x=tensor([0, 1]), y=tensor([0, 1]))
>>> # Example with `Tuple` inside the batch:
>>> default_collate([(0, 1), (2, 3)])
[tensor([0, 2]), tensor([1, 3])]
>>> # Example with `List` inside the batch:
>>> default_collate([[0, 1], [2, 3]])
[tensor([0, 2]), tensor([1, 3])]
>>> # Two options to extend `default_collate` to handle specific type
>>> # Option 1: Write custom collate function and invoke `default_collate`
>>> def custom_collate(batch):
...     elem = batch[0]
...     if isinstance(elem, CustomType):  # Some custom condition
...         return ...
...     else:  # Fall back to `default_collate`
...         return default_collate(batch)
>>> # Option 2: In-place modify `default_collate_fn_map`
>>> def collate_customtype_fn(batch, *, collate_fn_map=None):
...     return ...
>>> default_collate_fn_map.update(CustomType, collate_customtype_fn)
>>> default_collate(batch)  # Handle `CustomType` automatically

torch.utils.data.default_convert(data)[source]#

将每个 NumPy 数组元素转换为 torch.Tensor。

如果输入是 Sequence、Collection 或 Mapping，它会尝试将内部的每个元素转换为 torch.Tensor。如果输入不是 NumPy 数组，则保持不变。当 DataLoader 中未定义 batch_sampler 和 batch_size 时，此函数用作 collate 的默认函数。

通用输入类型到输出类型的映射与 default_collate() 类似。有关更多详细信息，请参阅那里的描述。

参数: data – 要转换的单个数据点

示例

>>> # Example with `int`
>>> default_convert(0)
0
>>> # Example with NumPy array
>>> default_convert(np.array([0, 1]))
tensor([0, 1])
>>> # Example with NamedTuple
>>> Point = namedtuple('Point', ['x', 'y'])
>>> default_convert(Point(0, 0))
Point(x=0, y=0)
>>> default_convert(Point(np.array(0), np.array(0)))
Point(x=tensor(0), y=tensor(0))
>>> # Example with List
>>> default_convert([np.array([0, 1]), np.array([2, 3])])
[tensor([0, 1]), tensor([2, 3])]

torch.utils.data.get_worker_info()[source]#

返回当前 DataLoader 迭代器工作进程的信息。

在工作进程中调用时，这将返回一个保证具有以下属性的对象：

id：当前工作进程 ID。
num_workers：总工作进程数。
seed：为当前工作进程设置的随机种子。此值由主进程 RNG 和工作进程 ID 确定。有关更多详细信息，请参阅 DataLoader 的文档。
dataset：此进程中 Dataset 对象的副本。请注意，这在不同进程中将是不同于主进程中对象的对象。

在主进程中调用时，返回 None。

注意

当在传递给 DataLoader 的 worker_init_fn 中使用时，此方法可用于不同地设置每个工作进程，例如，使用 worker_id 来配置 dataset 对象以仅读取分片 Dataset 的特定部分，或使用 seed 来为 Dataset 代码中使用的其他库设置种子。

返回类型: Optional[WorkerInfo]

torch.utils.data.random_split(dataset, lengths, generator=<torch._C.Generator object>)[source]#

将 Dataset 随机分割成给定长度的不重叠的新 Dataset。

如果提供了一个总和为 1 的分数列表，则长度将自动计算为每个提供分数的 floor(frac * len(dataset))。

在计算长度后，如果存在任何余数，则会将 1 的计数以循环方式分配给长度，直到没有余数为止。

可选地固定生成器以获得可复现的结果，例如：

示例

>>> generator1 = torch.Generator().manual_seed(42)
>>> generator2 = torch.Generator().manual_seed(42)
>>> random_split(range(10), [3, 7], generator=generator1)
>>> random_split(range(30), [0.3, 0.3, 0.4], generator=generator2)

参数

dataset (Dataset) – 要分割的 Dataset
lengths (sequence) – 要生成的分割的长度或分数
generator (Generator) – 用于随机排列的生成器。

返回类型

list[torch.utils.data.dataset.Subset[~_T]]

class torch.utils.data.Sampler(data_source=None)[source]#

所有 Sampler 的基类。

每个 Sampler 子类都必须提供一个 __iter__() 方法，提供一种迭代 Dataset 元素的索引或索引列表（批次）的方式，并且可以提供一个返回迭代器长度的 __len__() 方法。

参数: data_source (Dataset) – 此参数未使用，将在 2.2.0 中删除。您仍然可以有利用它的自定义实现。

示例

>>> class AccedingSequenceLengthSampler(Sampler[int]):
>>>     def __init__(self, data: List[str]) -> None:
>>>         self.data = data
>>>
>>>     def __len__(self) -> int:
>>>         return len(self.data)
>>>
>>>     def __iter__(self) -> Iterator[int]:
>>>         sizes = torch.tensor([len(x) for x in self.data])
>>>         yield from torch.argsort(sizes).tolist()
>>>
>>> class AccedingSequenceLengthBatchSampler(Sampler[List[int]]):
>>>     def __init__(self, data: List[str], batch_size: int) -> None:
>>>         self.data = data
>>>         self.batch_size = batch_size
>>>
>>>     def __len__(self) -> int:
>>>         return (len(self.data) + self.batch_size - 1) // self.batch_size
>>>
>>>     def __iter__(self) -> Iterator[List[int]]:
>>>         sizes = torch.tensor([len(x) for x in self.data])
>>>         for batch in torch.chunk(torch.argsort(sizes), len(self)):
>>>             yield batch.tolist()

注意

DataLoader 并不严格需要 __len__() 方法，但在任何涉及 DataLoader 长度的计算中都期望此方法。

class torch.utils.data.SequentialSampler(data_source)[source]#

顺序采样元素，始终按相同顺序。

参数: data_source (Dataset) – 要采样的 Dataset

class torch.utils.data.RandomSampler(data_source, replacement=False, num_samples=None, generator=None)[source]#

随机采样元素。如果不放回抽样，则从打乱的 Dataset 中采样。

如果放回抽样，则用户可以指定 num_samples 来抽取。

参数

data_source (Dataset) – 要采样的 Dataset
replacement (bool) – 如果为 True，则按需进行有放回抽样，默认为 ``False``
num_samples (int) – 要抽取的样本数量，默认为 `len(dataset)`。
generator (Generator) – 采样中使用的生成器。

class torch.utils.data.SubsetRandomSampler(indices, generator=None)[source]#

从给定的索引列表中随机采样元素，无放回。

参数

indices (sequence) – 索引序列
generator (Generator) – 采样中使用的生成器。

class torch.utils.data.WeightedRandomSampler(weights, num_samples, replacement=True, generator=None)[source]#

根据给定的概率（权重）从 [0,..,len(weights)-1] 中采样元素。

参数

weights (sequence) – 权重序列，不一定加起来等于一
num_samples (int) – 要抽取的样本数量
replacement (bool) – 如果为 True，则进行有放回抽样。如果为 False，则进行无放回抽样，这意味着当为一个行抽取一个样本索引时，该行不能再次抽取该索引。
generator (Generator) – 采样中使用的生成器。

示例

>>> list(WeightedRandomSampler([0.1, 0.9, 0.4, 0.7, 3.0, 0.6], 5, replacement=True))
[4, 4, 1, 4, 5]
>>> list(WeightedRandomSampler([0.9, 0.4, 0.05, 0.2, 0.3, 0.1], 5, replacement=False))
[0, 1, 4, 3, 2]

class torch.utils.data.BatchSampler(sampler, batch_size, drop_last)[source]#

包装另一个 sampler 以产生一个 mini-batch 的索引。

参数

sampler (Sampler or Iterable) – 基本 sampler。可以是任何可迭代对象
batch_size (int) – mini-batch 的大小。
drop_last (bool) – 如果为 True，则当最后一个批次的 size 小于 batch_size 时，sampler 将丢弃该批次。

示例

>>> list(BatchSampler(SequentialSampler(range(10)), batch_size=3, drop_last=False))
[[0, 1, 2], [3, 4, 5], [6, 7, 8], [9]]
>>> list(BatchSampler(SequentialSampler(range(10)), batch_size=3, drop_last=True))
[[0, 1, 2], [3, 4, 5], [6, 7, 8]]

class torch.utils.data.distributed.DistributedSampler(dataset, num_replicas=None, rank=None, shuffle=True, seed=0, drop_last=False)[source]#

将数据加载限制为 Dataset 的子集的 Sampler。

它在与 torch.nn.parallel.DistributedDataParallel 结合使用时特别有用。在这种情况下，每个进程都可以将 DistributedSampler 实例作为 DataLoader 的 sampler 传递，并加载仅属于该进程的原始 Dataset 的子集。

注意

假定 Dataset 的大小是恒定的，并且它的任何实例始终以相同的顺序返回相同的元素。

参数

dataset (Dataset) – 用于采样的 Dataset。
num_replicas (int, optional) – 参与分布式训练的进程数。默认情况下，world_size 从当前分布式组检索。
rank (int, optional) – 当前进程在 num_replicas 中的排名。默认情况下，rank 从当前分布式组检索。
shuffle (bool, optional) – 如果为 True（默认），sampler 将打乱索引。
seed (int, optional) – 如果 shuffle=True，则用于打乱 sampler 的随机种子。此数字在所有分布式组进程中应相同。默认：0。
drop_last (bool, optional) – 如果为 True，则 sampler 将丢弃数据的尾部，使其能够被复制数整除。如果为 False，则 sampler 将添加额外的索引以使数据能够被复制数整除。默认：False。

警告

在分布式模式下，在每个 epoch 开始时，**在**创建 DataLoader 迭代器**之前**调用 set_epoch() 方法对于正确地在多个 epoch 中进行打乱是必要的。否则，将始终使用相同的顺序。

示例

>>> sampler = DistributedSampler(dataset) if is_distributed else None
>>> loader = DataLoader(dataset, shuffle=(sampler is None),
...                     sampler=sampler)
>>> for epoch in range(start_epoch, n_epochs):
...     if is_distributed:
...         sampler.set_epoch(epoch)
...     train(loader)

torch.utils.data#

数据集类型#

Map-style 数据集#

Iterable-style 数据集#

数据加载顺序与 Sampler#

加载批处理和非批处理数据#

自动批处理（默认）#

禁用自动批处理#

使用 collate_fn#

单进程和多进程数据加载#

单进程数据加载（默认）#

多进程数据加载#

平台特定的行为#

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

内存固定#

文档

教程

资源

数据加载顺序与 `Sampler`#

使用 `collate_fn`#