DatasetとIterableDatasetの違い

データセットオブジェクトには、DatasetとIterableDatasetの2つの種類があります。 使用するまたは作成するデータセットの種類は、データセットのサイズによって決まります。 一般的に、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])

ただし、1つの注意点として、データセット全体をディスクまたはメモリに保存しておく必要があるため、ディスクより大きいデータセットにはアクセスできません。 大きなデータセットでは不便になる可能性があるため、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は、ディスクにファイルを書き込まずにオンラインデータを読み取ることができます。 例えば、C4やLAION-2Bのように、数百GBのシャードで構成されたデータセットをストリーミングできます。 データセットのストリーミング方法の詳細については、Dataset Streaming Guideを参照してください。

ただし、これが唯一の違いではありません。IterableDatasetの「lazy」な挙動は、データセットの作成や処理においても存在します。

マップスタイルのデータセットとイテラブルデータセットの作成

リストや辞書を使用して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 を用いて一度に1つのサンプルを返すため、通常の 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

ローカルファイルを完全に、かつ段階的に読み込む

ローカルまたはリモートのデータファイルをArrow Datasetに変換するには、load_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などの多くのファイル形式、および画像ファイルや音声ファイルがサポートされています。 tabular、text、vision、audioデータセットの読み込みに関する対応するガイドで、詳細を確認できます。

Eagerなデータ処理とlazyなデータ処理

DatasetオブジェクトをDataset.map()で処理すると、データセット全体が即座に処理され、結果が返されます。 これは例えば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をDataset.shuffle()でシャッフルすると、データセットが正確にシャッフルされます。 これは、インデックスのリスト[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() は、データセットが複数のファイルやソースで構成されている場合、データセットのシャードもシャッフルします。さらに、一度に最大 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

データセットのファイル数が少ない場合、シャッフルの前にreshard()を呼び出すことで、シャッフルの品質を向上させることができます。例えば、Parquetデータセットの場合、1つのシャードに1つのファイルという構成ではなく、行グループ単位でリシャーディングを行うため、シャッフル元のシャード数が大幅に増加します。

速度の違い

通常の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

データセットを複数のエポックで使用する場合、シャッフルバッファ内のシャード順序をシャッフルするための実効シードはseed + 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

マップスタイルのデータセットのイテレーションを再開するには、最初のサンプルをスキップするだけです:

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

ただし、DataLoaderにSamplerを使用する場合は、代わりにサンプラーの状態を保存する必要があります（再開を許可するカスタムサンプラーを作成している可能性があります）。

一方、iterable datasetは、再開するための特定の例のインデックスへのランダムアクセスを提供しません。ただし、モデルやオプティマイザーと同様に、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 datasetは現在読み込まれているshardと、そのshard内のexampleのインデックスを追跡し、この情報をstate_dictに保存します。

チェックポイントから再開する場合、データセットは以前に読み取られたすべてのシャードをスキップして、現在のシャードから再開します。 その後、シャードを読み取り、チェックポイントの正確なサンプルに到達するまでサンプルをスキップします。

したがって、すでに反復処理済みのシャードを再読み込みしないため、データセットの再起動は比較的速くなります。ただし、データセットの再開は通常、即座には完了しません。現在のシャードの先頭から読み込みを再開し、チェックポイントの位置に到達するまでサンプルをスキップする必要があるためです。

これはtorchdataのStatefulDataLoaderと組み合わせて使用でき、PyTorch DataLoaderによるストリーミングを参照してください。

mapスタイルからiterableスタイルへ切り替える

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