Datasets 文档
Dataset 和 IterableDataset 之间的差异
并获得增强的文档体验
开始使用
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
流式处理可以读取在线数据,而无需将任何文件写入磁盘。例如,您可以流式传输由多个分片组成的数据集,每个分片都有数百 GB,例如 C4 或 LAION-2B。在 数据集流式处理指南 中了解有关如何流式传输数据集的更多信息。
但这并不是唯一的区别,因为 IterableDataset
的“惰性”行为也存在于数据集创建和处理方面。
创建 map-style 数据集和 iterable 数据集
您可以使用列表或字典创建 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,以及图像和音频文件。您可以在加载表格、文本、视觉和 音频数据集的相应指南中找到更多信息。
立即数据处理和惰性数据处理
当您使用 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
但是,仅使用混洗缓冲区不足以提供令人满意的机器学习模型训练混洗。因此,如果您的数据集由多个文件或源组成,IterableDataset.shuffle() 也会混洗数据集分片
# 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
速度差异
常规 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 上使用数据集,则混洗缓冲区中混洗分片顺序的有效种子为 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-style 数据集的迭代,您可以简单地跳过前几个示例
my_dataset = my_dataset.select(range(start_index, len(dataset)))
但是,如果您将 DataLoader
与 Sampler
一起使用,则应改为保存采样器的状态(您可能已编写允许恢复的自定义采样器)。
另一方面,iterable 数据集不提供对特定示例索引的随机访问以进行恢复。但是,您可以使用 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)
在幕后,iterable 数据集会跟踪正在读取的当前分片以及当前分片中的示例索引,并将此信息存储在 state_dict
中。
要从检查点恢复,数据集会跳过先前读取的所有分片,以从当前分片重新开始。然后,它读取分片并跳过示例,直到到达检查点中的确切示例。
因此,重新启动数据集非常快,因为它不会重新读取已迭代的分片。尽管如此,恢复数据集通常不是瞬时的,因为它必须从当前分片的开头重新开始读取并跳过示例,直到到达检查点位置。
这可以与 torchdata
中的 StatefulDataLoader
一起使用,请参阅使用 PyTorch DataLoader 进行流式处理。
从 map-style 切换到 iterable
如果您想受益于 IterableDataset 的“惰性”行为或其速度优势,您可以将 map-style 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