Dataset 与 IterableDataset 之间的区别

数据集对象有两种类型：Dataset 和 IterableDataset。选择使用或创建哪种类型的数据集，取决于数据集的大小。通常，IterableDataset 因其惰性行为和速度优势，非常适合大型数据集（想想数百 GB 的规模！），而 Dataset 则适用于其他所有情况。本页将比较 Dataset 和 IterableDataset 之间的差异，帮助你选择合适的数据集对象。

下载和流式传输

当您拥有一个常规的 Dataset 时，可以使用 my_dataset[0] 来访问它。这提供了对行的随机访问。此类数据集也称为“map-style”数据集。例如，您可以像这样下载 ImageNet-1k 并访问任意一行：

from datasets import load_dataset

imagenet = load_dataset("timm/imagenet-1k-wds", split="train")  # downloads the full dataset
print(imagenet[0])

但有一个注意事项：您必须将整个数据集存储在磁盘或内存中，这会阻止您访问比磁盘更大的数据集。由于对于大型数据集来说这可能变得不方便，因此存在另一种类型的数据集，即 IterableDataset。当您拥有一个 IterableDataset 时，您可以使用 for 循环来访问它，在遍历数据集时逐步加载数据。这样，内存中只加载一小部分示例，并且不会在磁盘上写入任何内容。

例如，您可以在不将其下载到磁盘的情况下流式传输 ImageNet-1k 数据集：

from datasets import load_dataset

imagenet = load_dataset("timm/imagenet-1k-wds", split="train", streaming=True)  # will start loading the data when iterated over
for example in imagenet:
    print(example)
    break

Streaming 可以在不将任何文件写入磁盘的情况下读取在线数据。例如，你可以流式读取由多个分片组成的数据集，每个分片都有数百 GB 之大，例如 C4 或 LAION-2B。有关如何流式读取数据集的更多信息，请参阅 Dataset Streaming Guide。

不过，这并非唯一的区别，因为IterableDataset的“惰性”行为在数据集的创建和处理过程中同样存在。

创建 map 风格数据集和可迭代数据集

您可以使用列表或字典来创建 Dataset，数据会完全转换为 Arrow 格式，因此您可以轻松访问任意一行：

my_dataset = Dataset.from_dict({"col_1": [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]})
print(my_dataset[0])

另一方面，要创建一个 IterableDataset，你必须提供一种“惰性”的数据加载方式。在 Python 中，我们通常使用生成器函数。这些函数通过 yield 一次返回一个样本，这意味着你无法像访问常规 Dataset 那样通过切片来访问某一行：

def my_generator(n):
    for i in range(n):
        yield {"col_1": i}

my_iterable_dataset = IterableDataset.from_generator(my_generator, gen_kwargs={"n": 10})
for example in my_iterable_dataset:
    print(example)
    break

完全加载和渐进式加载本地文件

可以使用 load_dataset() 将本地或远程数据文件转换为 Arrow Dataset：

data_files = {"train": ["path/to/data.csv"]}
my_dataset = load_dataset("csv", data_files=data_files, split="train")
print(my_dataset[0])

然而，这需要从 CSV 到 Arrow 格式的转换步骤，如果您的数据集很大，这将耗费时间和磁盘空间。

为了节省磁盘空间并跳过转换步骤，您可以直接通过流式读取本地文件来定义一个 IterableDataset。这样，在遍历数据集时，数据会从本地文件中逐步读取：

data_files = {"train": ["path/to/data.csv"]}
my_iterable_dataset = load_dataset("csv", data_files=data_files, split="train", streaming=True)
for example in my_iterable_dataset:  # this reads the CSV file progressively as you iterate over the dataset
    print(example)
    break

支持多种文件格式，例如 CSV、JSONL 和 Parquet，以及图像和音频文件。您可以在加载表格、文本、视觉和音频数据集的相应指南中了解更多信息。

Eager 数据处理和 lazy 数据处理

使用 Dataset.map() 处理 Dataset 对象时，整个数据集会立即被处理并返回。这与例如 pandas 的工作方式类似。

my_dataset = my_dataset.map(process_fn)  # process_fn is applied on all the examples of the dataset
print(my_dataset[0])

另一方面，由于 IterableDataset 具有“惰性”特性，调用 IterableDataset.map() 并不会在整个数据集上应用你的 map 函数。相反，你的 map 函数是即时应用的。

因此，您可以串联多个处理步骤，在开始遍历数据集时，它们将同时运行：

my_iterable_dataset = my_iterable_dataset.map(process_fn_1)
my_iterable_dataset = my_iterable_dataset.filter(filter_fn)
my_iterable_dataset = my_iterable_dataset.map(process_fn_2)

# process_fn_1, filter_fn and process_fn_2 are applied on-the-fly when iterating over the dataset
for example in my_iterable_dataset:  
    print(example)
    break

精确且快速的近似洗牌

使用 Dataset.shuffle() 对 Dataset 进行洗牌时，会对数据集执行精确的洗牌操作。其工作原理是获取一个索引列表 [0, 1, 2, ... len(my_dataset) - 1]，然后对该列表进行洗牌。随后，访问 my_dataset[0] 将返回由洗牌后的索引映射中第一个元素所定义的该行及其索引：

my_dataset = my_dataset.shuffle(seed=42)
print(my_dataset[0])

由于 IterableDataset 无法对行进行随机访问，因此我们不能使用打乱顺序的索引列表来访问任意位置的行。这导致无法使用精确的洗牌方式。相反，IterableDataset.shuffle() 中使用了快速的近似洗牌方法。它利用洗牌缓冲区，从数据集中迭代地采样随机样本。由于数据集仍然是迭代读取的，因此它具有出色的速度性能：

my_iterable_dataset = my_iterable_dataset.shuffle(seed=42, buffer_size=100)
for example in my_iterable_dataset:
    print(example)
    break

但仅使用 shuffle buffer 不足以在机器学习模型训练中提供令人满意的打乱效果。因此，如果您的数据集由多个文件或数据源组成，IterableDataset.shuffle() 还会对数据集分片进行打乱。此外，它一次最多使用 max_buffer_input_shards 个分片来填充缓冲区（默认值为 10），这极大地提升了打乱的质量：

# Stream from the internet
my_iterable_dataset = load_dataset("deepmind/code_contests", split="train", streaming=True)
my_iterable_dataset.num_shards  # 39

# Stream from local files
data_files = {"train": [f"path/to/data_{i}.csv" for i in range(1024)]}
my_iterable_dataset = load_dataset("csv", data_files=data_files, split="train", streaming=True)
my_iterable_dataset.num_shards  # 1024

# From a generator function
def my_generator(n, sources):
    for source in sources:
        for example_id_for_current_source in range(n):
            yield {"example_id": f"{source}_{example_id_for_current_source}"}

gen_kwargs = {"n": 10, "sources": [f"path/to/data_{i}" for i in range(1024)]}
my_iterable_dataset = IterableDataset.from_generator(my_generator, gen_kwargs=gen_kwargs)
my_iterable_dataset.num_shards  # 1024

如果您的数据集文件较少，可以在执行 shuffle 之前调用 reshard() 来提升 shuffle 的质量。例如，对于 Parquet 数据集，此操作会使用 row group 进行重新分片，而不是每个 shard 对应一个文件，从而为您提供更多可用于 shuffle 的 shard。

速度差异

常规的 Dataset 对象基于 Arrow，后者提供了对行的快速随机访问。得益于内存映射以及 Arrow 是一种内存格式这一事实，从磁盘读取数据不会执行昂贵的系统调用和反序列化。当使用 for 循环进行迭代时，它通过迭代连续的 Arrow 记录批次，提供了更快的数据加载速度。

然而，一旦你的 Dataset 具有索引映射（例如通过 Dataset.shuffle()），速度可能会降低 10 倍。这是因为使用索引映射获取要读取的行索引需要额外的步骤，而且最关键的是，你不再读取连续的数据块。要恢复速度，你需要使用 Dataset.flatten_indices() 将整个数据集重新写入磁盘，这会移除索引映射。不过，根据数据集的大小，这可能需要大量时间：

my_dataset[0]  # fast
my_dataset = my_dataset.shuffle(seed=42)
my_dataset[0]  # up to 10x slower
my_dataset = my_dataset.flatten_indices()  # rewrite the shuffled dataset on disk as contiguous chunks of data
my_dataset[0]  # fast again

在这种情况下，我们建议切换到 IterableDataset 并利用其快速的近似洗牌方法 IterableDataset.shuffle()。它仅对分片顺序进行洗牌，并为您的数据集添加一个洗牌缓冲区，从而保持数据集的速度处于最优状态。您还可以轻松地对数据集重新洗牌：

for example in enumerate(my_iterable_dataset):  # fast
    pass

shuffled_iterable_dataset = my_iterable_dataset.shuffle(seed=42, buffer_size=100)

for example in enumerate(shuffled_iterable_dataset):  # as fast as before
    pass

shuffled_iterable_dataset = my_iterable_dataset.shuffle(seed=1337, buffer_size=100)  # reshuffling using another seed is instantaneous

for example in enumerate(shuffled_iterable_dataset):  # still as fast as before
    pass

如果您在多个 epoch 上使用数据集，用于打乱 shuffle buffer 中 shard 顺序的有效种子为 seed + epoch。这使得在 epoch 之间重新打乱数据集变得轻而易举：

for epoch in range(n_epochs):
    my_iterable_dataset.set_epoch(epoch)
    for example in my_iterable_dataset:  # fast + reshuffled at each epoch using `effective_seed = seed + epoch`
        pass

要重新开始迭代 map 风格数据集，只需跳过前几个样本即可：

my_dataset = my_dataset.select(range(start_index, len(dataset)))

但如果你使用带有 Sampler 的 DataLoader，则应改为保存采样器的状态（你可能编写了一个支持恢复的自定义采样器）。

另一方面，可迭代数据集无法通过特定的示例索引进行随机访问以从中恢复。但你可以使用 IterableDataset.state_dict() 和 IterableDataset.load_state_dict() 从检查点恢复，这与针对模型和优化器所做的操作类似：

>>> iterable_dataset = Dataset.from_dict({"a": range(6)}).to_iterable_dataset(num_shards=3)
>>> # save in the middle of training
>>> state_dict = iterable_dataset.state_dict()
>>> # and resume later
>>> iterable_dataset.load_state_dict(state_dict)

在底层，可迭代数据集会跟踪当前正在读取的分片以及当前分片中的示例索引，并将这些信息存储在 state_dict 中。

要从检查点恢复，数据集会跳过之前已读取的所有分片，从当前分片重新开始。然后它会读取该分片，并跳过示例，直到到达检查点中对应的确切示例。

因此，重启数据集的速度相当快，因为它不会重新读取已经迭代过的分片。不过，恢复数据集通常并非瞬间完成，因为它需要从当前分片的开头重新开始读取，并跳过示例，直到到达检查点位置。

这可以与 torchdata 中的 StatefulDataLoader 一起使用，请参阅使用 PyTorch DataLoader 进行流式处理。

从 map 样式切换到可迭代

如果你想利用 IterableDataset 的“惰性”行为或其速度优势，可以将你的 map 风格 Dataset 切换为 IterableDataset：

my_iterable_dataset = my_dataset.to_iterable_dataset()

如果您想要对数据集进行洗牌或将其与 PyTorch DataLoader 配合使用，我们建议生成一个分片的IterableDataset：

my_iterable_dataset = my_dataset.to_iterable_dataset(num_shards=1024)
my_iterable_dataset.num_shards  # 1024

Update on GitHub

Datasets

Dataset 与 IterableDataset 之间的区别

下载和流式传输

创建 map 风格数据集和可迭代数据集

完全加载和渐进式加载本地文件

Eager 数据处理和 lazy 数据处理

精确且快速的近似洗牌

速度差异

从 map 样式切换到可迭代