# torch.utils.data

> 原文： [https://pytorch.org/docs/stable/data.html](https://pytorch.org/docs/stable/data.html)

PyTorch 数据加载实用程序的核心是 [`torch.utils.data.DataLoader`](#torch.utils.data.DataLoader "torch.utils.data.DataLoader") 类。 它表示可在数据集上迭代的 Python，并支持

*   [地图样式和可迭代样式的数据集](#dataset-types)，

*   [自定义数据加载顺序](#data-loading-order-and-sampler)，

*   [自动配料](#loading-batched-and-non-batched-data)，

*   [单进程和多进程数据加载](#single-and-multi-process-data-loading)，

*   [自动内存固定](#memory-pinning)。

这些选项由 [`DataLoader`](#torch.utils.data.DataLoader "torch.utils.data.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)

```

以下各节详细介绍了这些选项的效果和用法。

## 数据集类型

[`DataLoader`](#torch.utils.data.DataLoader "torch.utils.data.DataLoader") 构造函数的最重要参数是`dataset`，它指示要从中加载数据的数据集对象。 PyTorch 支持两种不同类型的数据集：

*   [地图样式数据集](#map-style-datasets)，

*   [迭代式数据集](#iterable-style-datasets)。

### 地图样式数据集

映射样式数据集是一种实现`__getitem__()`和`__len__()`协议的数据集，它表示从(可能是非整数）索引/关键字到数据样本的映射。

例如，当使用`dataset[idx]`访问时，此类数据集可以从磁盘上的文件夹中读取第`idx`张图像及其对应的标签。

有关更多详细信息，请参见 [`Dataset`](#torch.utils.data.Dataset "torch.utils.data.Dataset") 。

### 迭代式数据集

可迭代样式的数据集是 [`IterableDataset`](#torch.utils.data.IterableDataset "torch.utils.data.IterableDataset") 子类的实例，该子类实现了`__iter__()`协议，并表示数据样本上的可迭代。 这种类型的数据集特别适用于随机读取价格昂贵甚至不大可能，并且批处理大小取决于所获取数据的情况。

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

有关更多详细信息，请参见 [`IterableDataset`](#torch.utils.data.IterableDataset "torch.utils.data.IterableDataset") 。

注意

当将 [`IterableDataset`](#torch.utils.data.IterableDataset "torch.utils.data.IterableDataset") 与[一起使用时，多进程数据加载](#multi-process-data-loading)。 在每个工作进程上都复制相同的数据集对象，因此必须对副本进行不同的配置，以避免重复的数据。 有关如何实现此功能的信息，请参见 [`IterableDataset`](#torch.utils.data.IterableDataset "torch.utils.data.IterableDataset") 文档。

## 数据加载顺序和 [`Sampler`](#torch.utils.data.Sampler "torch.utils.data.Sampler")

对于[迭代式数据集](#iterable-style-datasets)，数据加载顺序完全由用户定义的迭代器控制。 这样可以更轻松地实现块读取和动态批次大小的实现(例如，通过每次生成一个批次的样本）。

本节的其余部分涉及[地图样式数据集](#map-style-datasets)的情况。 [`torch.utils.data.Sampler`](#torch.utils.data.Sampler "torch.utils.data.Sampler") 类用于指定数据加载中使用的索引/键的顺序。 它们代表数据集索引上的可迭代对象。 例如，在具有随机梯度体面(SGD）的常见情况下， [`Sampler`](#torch.utils.data.Sampler "torch.utils.data.Sampler") 可以随机排列一列索引，一次生成每个索引，或者为小批量生成少量索引 新币。

基于 [`DataLoader`](#torch.utils.data.DataLoader "torch.utils.data.DataLoader") 的`shuffle`参数，将自动构建顺序采样或混洗的采样器。 或者，用户可以使用`sampler`参数指定一个自定义 [`Sampler`](#torch.utils.data.Sampler "torch.utils.data.Sampler") 对象，该对象每次都会产生要提取的下一个索引/关键字。

可以一次生成批量索引列表的自定义 [`Sampler`](#torch.utils.data.Sampler "torch.utils.data.Sampler") 作为`batch_sampler`参数传递。 也可以通过`batch_size`和`drop_last`参数启用自动批处理。 有关更多详细信息，请参见下一部分的[。](#loading-batched-and-non-batched-data)

Note

`sampler`和`batch_sampler`都不与可迭代样式的数据集兼容，因为此类数据集没有键或索引的概念。

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

[`DataLoader`](#torch.utils.data.DataLoader "torch.utils.data.DataLoader") 支持通过参数`batch_size`，`drop_last`和`batch_sampler`将各个提取的数据样本自动整理为批次。

### 自动批处理(默认）

这是最常见的情况，对应于获取一小批数据并将其整理为批处理的样本，即包含张量，其中一维为批处理维度(通常是第一维）。

当`batch_size`(默认`1`）不是`None`时，数据加载器将生成批处理的样本，而不是单个样本。 `batch_size`和`drop_last`参数用于指定数据加载器如何获取数据集密钥的批处理。 对于地图样式的数据集，用户可以选择指定`batch_sampler`，它一次生成一个键列表。

Note

`batch_size`和`drop_last`自变量本质上用于从`sampler`构造`batch_sampler`。 对于地图样式的数据集，`sampler`由用户提供或基于`shuffle`参数构造。 对于可迭代样式的数据集，`sampler`是一个虚拟的无限数据集。 有关采样器的更多详细信息，请参见[本节](#data-loading-order-and-sampler)。

Note

当从[可重复样式数据集](#iterable-style-datasets)进行[多重处理](#multi-process-data-loading)提取时，`drop_last`参数会删除每个工作人员数据集副本的最后一个非完整批次。

使用来自采样器的索引获取样本列表后，作为`collate_fn`参数传递的函数用于将样本列表整理为批次。

在这种情况下，从地图样式数据集加载大致等效于：

```
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`的更多信息，请参见[本部分](#dataloader-collate-fn)。

### 禁用自动批处理

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

当`batch_size`和`batch_sampler`均为`None`时(`batch_sampler`的默认值已为`None`），自动批处理被禁用。 从`dataset`获得的每个样本都将作为`collate_fn`参数传递的函数进行处理。

**禁用自动批处理**时，默认值`collate_fn`仅将 NumPy 数组转换为 PyTorch 张量，而其他所有内容均保持不变。

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)

```

有关`collate_fn`的更多信息，请参见[本部分](#dataloader-collate-fn)。

### 使用`collate_fn`

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

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

**启用自动批处理**时，会每次调用`collate_fn`并带有数据样本列表。 期望将输入样本整理为一批，以便从数据加载器迭代器中获得收益。 本节的其余部分描述了这种情况下默认`collate_fn`的行为。

例如，如果每个数据样本都包含一个 3 通道图像和一个整体类标签，即数据集的每个元素返回一个元组`(image, class_index)`，则默认值`collate_fn`将此类元组的列表整理为一个元组 批处理图像张量和批处理类标签 Tensor。 特别是，默认`collate_fn`具有以下属性：

*   它始终将新维度添加为批次维度。

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

*   它保留了数据结构，例如，如果每个样本都是一个字典，它将输出一个具有相同键集但将批处理 Tensors 作为值的字典(如果无法将这些值转换为 Tensors，则将其列出）。 与`list`，`tuple`，`namedtuple`等相同。

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

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

默认情况下， [`DataLoader`](#torch.utils.data.DataLoader "torch.utils.data.DataLoader") 使用单进程数据加载。

在 Python 进程中，[全局解释器锁定(GIL）](https://wiki.python.org/moin/GlobalInterpreterLock)阻止了跨线程真正的完全并行化 Python 代码。 为了避免在加载数据时阻塞计算代码，PyTorch 提供了一个简单的开关，只需将参数`num_workers`设置为正整数即可执行多进程数据加载。

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

在此模式下，以与初始化 [`DataLoader`](#torch.utils.data.DataLoader "torch.utils.data.DataLoader") 相同的过程完成数据提取。 因此，数据加载可能会阻止计算。 然而，当用于在进程之间共享数据的资源(例如，共享存储器，文件描述符）受到限制时，或者当整个数据集很小并且可以完全加载到存储器中时，该模式可能是优选的。 此外，单进程加载通常显示更多可读的错误跟踪，因此对于调试很有用。

### 多进程数据加载

将参数`num_workers`设置为正整数将打开具有指定数量的加载程序工作进程的多进程数据加载。

在此模式下，每次创建 [`DataLoader`](#torch.utils.data.DataLoader "torch.utils.data.DataLoader") 的迭代器时(例如，当您调用`enumerate(dataloader)`时），都会创建`num_workers`工作进程。 此时，`dataset`，`collate_fn`和`worker_init_fn`被传递给每个工作程序，在这里它们被用来初始化和获取数据。 这意味着数据集访问及其内部 IO 转换(包括`collate_fn`）在工作进程中运行。

[`torch.utils.data.get_worker_info()`](#torch.utils.data.get_worker_info "torch.utils.data.get_worker_info") 在工作进程中返回各种有用的信息(包括工作 ID，数据集副本，初始种子等），并在主进程中返回`None`。 用户可以在数据集代码和/或`worker_init_fn`中使用此功能来分别配置每个数据集副本，并确定代码是否正在工作进程中运行。 例如，这在分片数据集时特别有用。

对于地图样式的数据集，主过程使用`sampler`生成索引并将其发送给工作人员。 因此，任何随机播放都是在主过程中完成的，该过程通过为索引分配索引来引导加载。

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

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

警告

通常不建议在多进程加载中返回 CUDA 张量，因为在使用 CUDA 和在并行处理中共享 CUDA 张量时存在很多微妙之处(请参见[在并行处理中的 CUDA](notes/multiprocessing.html#multiprocessing-cuda-note))。 相反，我们建议使用[自动内存固定](#memory-pinning)(即，设置`pin_memory=True`），该功能可以将数据快速传输到支持 CUDA 的 GPU。

#### 平台特定的行为

由于工作程序依赖于 Python `multiprocessing`，因此与 Unix 相比，Windows 上的工作程序启动行为有所不同。

*   在 Unix 上，`fork()`是默认的`multiprocessing`启动方法。 使用`fork()`，童工通常可以直接通过克隆的地址空间访问`dataset`和 Python 参数函数。

*   在 Windows 上，`spawn()`是默认的`multiprocessing`启动方法。 使用`spawn()`启动另一个解释器，该解释器运行您的主脚本，然后运行内部工作程序函数，该函数通过序列化`pickle`接收`dataset`，`collate_fn`和其他参数。

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

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

*   确保在`__main__`检查之外将任何自定义`collate_fn`，`worker_init_fn`或`dataset`代码声明为顶级定义。 这样可以确保它们在工作进程中可用。 (这是必需的，因为将函数仅作为引用而不是`bytecode`进行腌制。）

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

默认情况下，每个工作人员的 PyTorch 种子将设置为`base_seed + worker_id`，其中`base_seed`是主进程使用其 RNG 生成的长整数(因此，强制使用 RNG 状态）。 但是，初始化工作程序(例如 NumPy）时，可能会复制其他库的种子，导致每个工作程序返回相同的随机数。 (请参阅 FAQ 中的[本部分](notes/faq.html#dataloader-workers-random-seed)。）

在`worker_init_fn`中，您可以使用 [`torch.utils.data.get_worker_info().seed`](#torch.utils.data.get_worker_info "torch.utils.data.get_worker_info") 或 [`torch.initial_seed()`](torch.html#torch.initial_seed "torch.initial_seed") 访问每个工作人员的 PyTorch 种子集，并在加载数据之前使用它为其他库提供种子。

## 内存固定

主机到 GPU 副本源自固定(页面锁定）内存时，速度要快得多。 有关通常何时以及如何使用固定内存的更多详细信息，请参见[使用固定内存缓冲区](notes/cuda.html#cuda-memory-pinning)。

对于数据加载，将`pin_memory=True`传递到 [`DataLoader`](#torch.utils.data.DataLoader "torch.utils.data.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=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)¶
```

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

[`DataLoader`](#torch.utils.data.DataLoader "torch.utils.data.DataLoader") 支持地图样式和可迭代样式的数据集，具有单进程或多进程加载，自定义加载顺序以及可选的自动批处理(归类）和内存固定。

有关更多详细信息，请参见 [`torch.utils.data`](#module-torch.utils.data "torch.utils.data") 文档页面。

参数

*   **数据集** ([_数据集_](#torch.utils.data.Dataset "torch.utils.data.Dataset"))–要从中加载数据的数据集。

*   **batch_size**  (_python：int_ _，_ _可选_）–每批次要加载多少个样本(默认值：`1`）。

*   **随机播放** (_bool_ _，_ _可选_）–设置为`True`以使数据在每个时间段都重新随机播放(默认值：`False` )。

*   **采样器** ([_采样器_](#torch.utils.data.Sampler "torch.utils.data.Sampler") _，_ _可选_）–定义了从数据集中抽取样本的策略。 如果指定，则`shuffle`必须为`False`。

*   **batch_sampler**  ([_采样器_](#torch.utils.data.Sampler "torch.utils.data.Sampler") _，_ _可选_）–类似`sampler`，但在 时间。 与`batch_size`，`shuffle`，`sampler`和`drop_last`互斥。

*   **num_workers**  (_python：int_ _，_ _可选_）–多少个子进程用于数据加载。 `0`表示将在主进程中加载​​数据。 (默认：`0`）

*   **collat​​e_fn** (可调用的_，_ _可选_）–合并样本列表以形成张量的小批量。 在从地图样式数据集中使用批量加载时使用。

*   **pin_memory**  (_bool_ _，_ _可选_）–如果`True`，则数据加载器将张量复制到 CUDA 固定的内存中，然后返回。 如果您的数据元素是自定义类型，或者您的`collate_fn`返回的是自定义类型的批次，请参见下面的示例。

*   **drop_last** (_布尔_ _，_ _可选_）–设置为`True`以删除最后不完整的批次，如果数据集大小不可分割 按批次大小。 如果`False`并且数据集的大小不能被批次大小整除，那么最后一批将较小。 (默认：`False`）

*   **超时**(_数字_ _，_ _可选_）–如果为正，则表示从工作人员处收集批次的超时值。 应始终为非负数。 (默认：`0`）

*   **worker_init_fn** (_可调用_ _，_ _可选_）–如果不是`None`，则将在每个具有工作人员 ID (在播种之后和数据加载之前，将`[0, num_workers - 1]`中的 int 作为输入。 (默认：`None`）

Warning

如果使用`spawn`启动方法，则`worker_init_fn`不能是不可拾取的对象，例如 lambda 函数。 有关 PyTorch 中与并行处理有关的更多详细信息，请参见[并行处理最佳实践](notes/multiprocessing.html#multiprocessing-best-practices)。

Note

`len(dataloader)`启发式方法基于所用采样器的长度。 当`dataset`是 [`IterableDataset`](#torch.utils.data.IterableDataset "torch.utils.data.IterableDataset") 时，无论多进程加载配置如何，都将返回`len(dataset)`(如果实现），因为 PyTorch 信任用户`dataset`代码可以正确处理多进程加载 避免重复数据。 有关这两种类型的数据集以及 [`IterableDataset`](#torch.utils.data.IterableDataset "torch.utils.data.IterableDataset") 如何与[多进程数据加载](#multi-process-data-loading)交互的更多详细信息，请参见[数据集类型](#dataset-types)。

* * *

```
class torch.utils.data.Dataset¶
```

表示 [`Dataset`](#torch.utils.data.Dataset "torch.utils.data.Dataset") 的抽象类。

代表从键到数据样本的映射的所有数据集都应将其子类化。 所有子类都应该覆盖`__getitem__()`，支持为给定键获取数据样本。 子类还可以选择覆盖`__len__()`，它有望通过许多 [`Sampler`](#torch.utils.data.Sampler "torch.utils.data.Sampler") 实现以及 [`DataLoader`](#torch.utils.data.DataLoader "torch.utils.data.DataLoader") 的默认选项返回数据集的大小。

Note

默认情况下， [`DataLoader`](#torch.utils.data.DataLoader "torch.utils.data.DataLoader") 构造一个索引采样器，该采样器产生整数索引。 要使其与具有非整数索引/键的地图样式数据集一起使用，必须提供自定义采样器。

* * *

```
class torch.utils.data.IterableDataset¶
```

可迭代的数据集。

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

所有子类都应覆盖`__iter__()`，这将返回此数据集中的样本迭代器。

当子类与 [`DataLoader`](#torch.utils.data.DataLoader "torch.utils.data.DataLoader") 一起使用时，数据集中的每个项目都将由 [`DataLoader`](#torch.utils.data.DataLoader "torch.utils.data.DataLoader") 迭代器产生。 当`num_workers &gt; 0`时，每个工作进程将具有数据集对象的不同副本，因此通常需要独立配置每个副本，以避免从工作进程返回重复的数据。 [`get_worker_info()`](#torch.utils.data.get_worker_info "torch.utils.data.get_worker_info") 在工作程序进程中调用时，返回有关工作程序的信息。 可以在数据集的`__iter__()`方法或 [`DataLoader`](#torch.utils.data.DataLoader "torch.utils.data.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)))
[3, 4, 5, 6]

>>> # Mult-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)))
[3, 5, 4, 6]

>>> # With even more workers
>>> print(list(torch.utils.data.DataLoader(ds, num_workers=20)))
[3, 4, 5, 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=20, worker_init_fn=worker_init_fn)))
[3, 4, 5, 6]

```

* * *

```
class torch.utils.data.TensorDataset(*tensors)¶
```

数据集包装张量。

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

Parameters

***张量** ([_tensor_](tensors.html#torch.Tensor "torch.Tensor"))–具有与第一维相同大小的张量。

* * *

```
class torch.utils.data.ConcatDataset(datasets)¶
```

数据集是多个数据集的串联。

此类对于组装不同的现有数据集很有用。

Parameters

**数据集**(_序列_）–要连接的数据集列表

* * *

```
class torch.utils.data.ChainDataset(datasets)¶
```

用于链接多个 [`IterableDataset`](#torch.utils.data.IterableDataset "torch.utils.data.IterableDataset") 的数据集。

此类对于组装不同的现有数据集流很有用。 链接操作是即时完成的，因此将大型数据集与此类连接起来将非常有效。

Parameters

**数据集**(IterableDataset 的_可迭代）–链接在一起的数据集_

* * *

```
class torch.utils.data.Subset(dataset, indices)¶
```

指定索引处的数据集子集。

Parameters

*   **数据集** ([_数据集_](#torch.utils.data.Dataset "torch.utils.data.Dataset"))–整个数据集

*   **索引**(_序列_）–为子集选择的整个集合中的索引

* * *

```
torch.utils.data.get_worker_info()¶
```

返回有关当前 [`DataLoader`](#torch.utils.data.DataLoader "torch.utils.data.DataLoader") 迭代器工作进程的信息。

在工作线程中调用时，此方法返回一个保证具有以下属性的对象：

*   `id`：当前工作人员 ID。

*   `num_workers`：工人总数。

*   `seed`：当前工作程序的随机种子集。 该值由主进程 RNG 和工作程序 ID 确定。 有关更多详细信息，请参见 [`DataLoader`](#torch.utils.data.DataLoader "torch.utils.data.DataLoader") 的文档。

*   `dataset`：此流程在**中的数据集对象的副本。 请注意，在不同的过程中，这将是与主过程中的对象不同的对象。**

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

Note

在传递给 [`DataLoader`](#torch.utils.data.DataLoader "torch.utils.data.DataLoader") 的`worker_init_fn`中使用时，此方法可用于不同地设置每个工作进程，例如，使用`worker_id`将`dataset`对象配置为仅读取 分片数据集的特定部分，或使用`seed`播种数据集代码中使用的其他库(例如 NumPy）。

* * *

```
torch.utils.data.random_split(dataset, lengths)¶
```

将数据集随机拆分为给定长度的不重叠的新数据集。

Parameters

*   **数据集** ([_数据集_](#torch.utils.data.Dataset "torch.utils.data.Dataset"))–要拆分的数据集

*   **长度**(_序列_）–要产生的分割的长度

* * *

```
class torch.utils.data.Sampler(data_source)¶
```

所有采样器的基类。

每个 Sampler 子类都必须提供`__iter__()`方法(提供一种对数据集元素的索引进行迭代的方法）和`__len__()`方法，该方法返回返回的迭代器的长度。

Note

[`DataLoader`](#torch.utils.data.DataLoader "torch.utils.data.DataLoader") 并非严格要求`__len__()`方法，但在涉及 [`DataLoader`](#torch.utils.data.DataLoader "torch.utils.data.DataLoader") 长度的任何计算中都应采用。

* * *

```
class torch.utils.data.SequentialSampler(data_source)¶
```

始终以相同顺序顺序采样元素。

Parameters

**data_source**  ([_数据集_](#torch.utils.data.Dataset "torch.utils.data.Dataset"))–要从中采样的数据集

* * *

```
class torch.utils.data.RandomSampler(data_source, replacement=False, num_samples=None)¶
```

随机采样元素。 如果不进行替换，则从经过改组的数据集中采样。 如果要更换，则用户可以指定`num_samples`进行绘制。

Parameters

*   **data_source** ([_Dataset_](#torch.utils.data.Dataset "torch.utils.data.Dataset")) – dataset to sample from

*   **替换** (_bool_ )–如果`True`为默认值，则替换为`True`

*   **num_samples**  (_python：int_ )–要绘制的样本数，默认为 len(dataset）。 仅当&lt;cite&gt;替换&lt;/cite&gt;为`True`时才应指定此参数。

* * *

```
class torch.utils.data.SubsetRandomSampler(indices)¶
```

从给定的索引列表中随机抽样元素，而无需替换。

Parameters

**索引**(_序列_）–索引序列

* * *

```
class torch.utils.data.WeightedRandomSampler(weights, num_samples, replacement=True)¶
```

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

Parameters

*   **权重**(_序列_）–权重序列，不必累加一个

*   **num_samples**  (_python：int_ )–要绘制的样本数

*   **替代品** (_bool_ )–如果`True`，则抽取替代品抽取样品。 如果没有，则它们将被替换而不会被绘制，这意味着当为一行绘制样本索引时，无法为该行再次绘制它。

例

```
>>> list(WeightedRandomSampler([0.1, 0.9, 0.4, 0.7, 3.0, 0.6], 5, replacement=True))
[0, 0, 0, 1, 0]
>>> 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)¶
```

包装另一个采样器以产生一个小批量的索引。

Parameters

*   **采样器** ([_采样器_](#torch.utils.data.Sampler "torch.utils.data.Sampler"))–基本采样器。

*   **batch_size**  (_python：int_ )–迷你批量的大小。

*   **drop_last**  (_bool_ )–如果为`True`，则采样器将丢弃最后一批，如果其大小小于`batch_size`

Example

```
>>> 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)¶
```

将数据加载限制为数据集子集的采样器。

与 [`torch.nn.parallel.DistributedDataParallel`](nn.html#torch.nn.parallel.DistributedDataParallel "torch.nn.parallel.DistributedDataParallel") 结合使用时特别有用。 在这种情况下，每个进程都可以将 DistributedSampler 实例作为 DataLoader 采样器传递，并加载原始数据集的专有子集。

Note

假定数据集大小恒定。

Parameters

*   **数据集** –用于采样的数据集。

*   **num_replicas** (_可选_）–参与分布式训练的进程数。

*   **等级**(_可选_）–当前进程在 num_replicas 中的等级。

*   **随机播放**(_可选_）–如果为 true(默认值），采样器将随机播放索引