Datasets documentation

PyTorchでの使用

Datasets

Join the Hugging Face community

and get access to the augmented documentation experience

Collaborate on models, datasets and Spaces

Faster examples with accelerated inference

Switch between documentation themes

to get started

PyTorchでの使用

このドキュメントは、datasetsをPyTorchと組み合わせて使用する際の簡単な入門ガイドです。特に、データセットからtorch.Tensorオブジェクトを取得する方法、およびPyTorchのDataLoaderとHugging FaceのDatasetを最高のパフォーマンスで使用する 방법에焦点を当てています。

データセットの形式

デフォルトでは、datasetsは通常のPythonオブジェクト（整数、浮動小数点数、文字列、リストなど）を返します。

代わりにPyTorchテンソルを取得するには、Dataset.with_format()を使用して、データセットの形式をpytorchに設定できます。

>>> from datasets import Dataset
>>> data = [[1, 2],[3, 4]]
>>> ds = Dataset.from_dict({"data": data})
>>> ds = ds.with_format("torch")
>>> ds[0]
{'data': tensor([1, 2])}
>>> ds[:2]
{'data': tensor([[1, 2],
         [3, 4]])}

Dataset オブジェクトは Arrow テーブルのラッパーであり、データセット内の配列から PyTorch テンソルへの高速なゼロコピー読み取りを可能にします。

GPU上でテンソルとしてデータを読み込むには、device引数を指定します:

>>> import torch
>>> device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
>>> ds = ds.with_format("torch", device=device)
>>> ds[0]
{'data': tensor([1, 2], device='cuda:0')}

N次元配列

データセットがN次元の配列で構成されている場合、デフォルトでは、shapeが固定されている場合、それらは同じtensorとして扱われることに注意してください。

>>> from datasets import Dataset
>>> data = [[[1, 2],[3, 4]],[[5, 6],[7, 8]]]  # fixed shape
>>> ds = Dataset.from_dict({"data": data})
>>> ds = ds.with_format("torch")
>>> ds[0]
{'data': tensor([[1, 2],
         [3, 4]])}

>>> from datasets import Dataset
>>> data = [[[1, 2],[3]],[[4, 5, 6],[7, 8]]]  # varying shape
>>> ds = Dataset.from_dict({"data": data})
>>> ds = ds.with_format("torch")
>>> ds[0]
{'data': [tensor([1, 2]), tensor([3])]}

ただし、このロジックは多くの場合、遅い形状の比較とデータのコピーを必要とします。これを回避するには、Array機能型を明示的に使用し、テンソルの形状を指定する必要があります：

>>> from datasets import Dataset, Features, Array2D
>>> data = [[[1, 2],[3, 4]],[[5, 6],[7, 8]]]
>>> features = Features({"data": Array2D(shape=(2, 2), dtype='int32')})
>>> ds = Dataset.from_dict({"data": data}, features=features)
>>> ds = ds.with_format("torch")
>>> ds[0]
{'data': tensor([[1, 2],
         [3, 4]])}
>>> ds[:2]
{'data': tensor([[[1, 2],
          [3, 4]],
 
         [[5, 6],
          [7, 8]]])}

その他の特徴タイプ

ClassLabel データは正しくテンソルに変換されます:

>>> from datasets import Dataset, Features, ClassLabel
>>> labels = [0, 0, 1]
>>> features = Features({"label": ClassLabel(names=["negative", "positive"])})
>>> ds = Dataset.from_dict({"label": labels}, features=features) 
>>> ds = ds.with_format("torch")  
>>> ds[:3]
{'label': tensor([0, 0, 1])}

PyTorchは数値のみをサポートするため、文字列オブジェクトとバイナリオブジェクトは変更されません。

ImageおよびAudioの機能タイプもサポートされています。

Image機能タイプを使用するには、visionエクストラをpip install datasets[vision]としてインストールする必要があります。

>>> from datasets import Dataset, Features, Audio, Image
>>> images = ["path/to/image.png"] * 10
>>> features = Features({"image": Image()})
>>> ds = Dataset.from_dict({"image": images}, features=features) 
>>> ds = ds.with_format("torch")
>>> ds[0]["image"].shape
torch.Size([512, 512, 4])
>>> ds[0]
{'image': tensor([[[255, 215, 106, 255],
         [255, 215, 106, 255],
         ...,
         [255, 255, 255, 255],
         [255, 255, 255, 255]]], dtype=torch.uint8)}
>>> ds[:2]["image"].shape
torch.Size([2, 512, 512, 4])
>>> ds[:2]
{'image': tensor([[[[255, 215, 106, 255],
          [255, 215, 106, 255],
          ...,
          [255, 255, 255, 255],
          [255, 255, 255, 255]]]], dtype=torch.uint8)}

Audio機能タイプを使用するには、audioエクストラをpip install datasets[audio]としてインストールする必要があります。

>>> from datasets import Dataset, Features, Audio, Image
>>> audio = ["path/to/audio.wav"] * 10
>>> features = Features({"audio": Audio()})
>>> ds = Dataset.from_dict({"audio": audio}, features=features) 
>>> ds = ds.with_format("torch")  
>>> ds[0]["audio"]["array"]
tensor([ 6.1035e-05,  1.5259e-05,  1.6785e-04,  ..., -1.5259e-05,
        -1.5259e-05,  1.5259e-05])
>>> ds[0]["audio"]["sampling_rate"]
tensor(44100)

データ読み込み

torch.utils.data.Dataset オブジェクトと同様に、Dataset は PyTorch の DataLoader に直接渡すことができます:

>>> import numpy as np
>>> from datasets import Dataset 
>>> from torch.utils.data import DataLoader
>>> data = np.random.rand(16)
>>> label = np.random.randint(0, 2, size=16)
>>> ds = Dataset.from_dict({"data": data, "label": label}).with_format("torch")
>>> dataloader = DataLoader(ds, batch_size=4)
>>> for batch in dataloader:
...     print(batch)                                                                                            
{'data': tensor([0.0047, 0.4979, 0.6726, 0.8105]), 'label': tensor([0, 1, 0, 1])}
{'data': tensor([0.4832, 0.2723, 0.4259, 0.2224]), 'label': tensor([0, 0, 0, 0])}
{'data': tensor([0.5837, 0.3444, 0.4658, 0.6417]), 'label': tensor([0, 1, 0, 0])}
{'data': tensor([0.7022, 0.1225, 0.7228, 0.8259]), 'label': tensor([1, 1, 1, 1])}

データ読み込みの最適化

データの読み込み速度を向上させる方法はいくつかあり、特に大規模なデータセットを扱う場合、時間の節約になります。 PyTorchでは、並列化されたデータ読み込み、個別にではなくインデックスのバッチを取得する方式、およびディスクにダウンロードせずにデータセットを反復処理するためのストリーミングが提供されています。

複数のワーカーを使用する

PyTorchのDataLoaderのnum_workers引数を使用してデータ読み込みを並列化し、スループットを向上させることができます。

内部では、DataLoaderがnum_workers個のプロセスを開始します。各プロセスは、DataLoaderに渡されたデータセットを再読み込みし、サンプルのクエリに使用されます。ワーカー内でデータセットを再読み込みしてもRAMが埋まることはなく、ディスクからデータセットをメモリマップし直すだけだからです。

>>> import numpy as np
>>> from datasets import Dataset, load_from_disk
>>> from torch.utils.data import DataLoader
>>> data = np.random.rand(10_000)
>>> Dataset.from_dict({"data": data}).save_to_disk("my_dataset")
>>> ds = load_from_disk("my_dataset").with_format("torch")
>>> dataloader = DataLoader(ds, batch_size=32, num_workers=4)

データをストリーミングする

IterableDataset として読み込むことで、データセットをストリーミングできます。これにより、リモートデータセットをディスクにダウンロードせずに、またはローカルのデータファイルに対して、段階的に反復処理を行うことができます。どのタイプのデータセットがユースケースに最適であるかについては、通常データセットとイテラブルデータセットの選択ガイドを参照してください。

datasetsのイテラブルなデータセットはtorch.utils.data.IterableDatasetを継承しているため、torch.utils.data.DataLoaderに渡すことができます。

>>> import numpy as np
>>> from datasets import Dataset, load_dataset
>>> from torch.utils.data import DataLoader
>>> data = np.random.rand(10_000)
>>> Dataset.from_dict({"data": data}).push_to_hub("<username>/my_dataset")  # Upload to the Hugging Face Hub
>>> my_iterable_dataset = load_dataset("<username>/my_dataset", streaming=True, split="train")
>>> dataloader = DataLoader(my_iterable_dataset, batch_size=32)

データセットが複数のシャードに分割されている場合（つまり、データセットが複数のデータファイルで構成されている場合）、num_workers を使用して並列ストリーミングを行うことができます。

>>> my_iterable_dataset = load_dataset("deepmind/code_contests", streaming=True, split="train")
>>> my_iterable_dataset.num_shards
39
>>> dataloader = DataLoader(my_iterable_dataset, batch_size=32, num_workers=4)

この場合、各ワーカーには、ストリーミングするシャードのリストのサブセットが割り当てられます。

チェックポイントと再開

トレーニングの途中でチェックポイントを作成して再開できるDataLoaderが必要な場合は、torchdataのStatefulDataLoaderを使用できます。

>>> from torchdata.stateful_dataloader import StatefulDataLoader
>>> my_iterable_dataset = load_dataset("deepmind/code_contests", streaming=True, split="train")
>>> dataloader = StatefulDataLoader(my_iterable_dataset, batch_size=32, num_workers=4)
>>> # save in the middle of training
>>> state_dict = dataloader.state_dict()
>>> # and resume later
>>> dataloader.load_state_dict(state_dict)

これは、IterableDataset.state_dict() と IterableDataset.load_state_dict() によって可能になります。

分散

To split your dataset across your training nodes, you can use datasets.distributed.split_dataset_by_node():

import os
from datasets.distributed import split_dataset_by_node

ds = split_dataset_by_node(ds, rank=int(os.environ["RANK"]), world_size=int(os.environ["WORLD_SIZE"]))

これは、mapスタイルのデータセットとiterableデータセットの両方で動作します。データセットは、world_sizeのノードからなるプール内の、ランクrankのノード用に分割されます。

マップスタイルのデータセットの場合:

各ノードにはデータのチャンクが割り当てられます。例えば、rank 0 にはデータセットの最初のチャンクが割り当てられます。

イテラブルなデータセットの場合:

strategy の分割は、"auto"、"shards"、または "examples" のいずれかです:

"shards": シャードはノード間で均等に割り当てられます
"examples": 各ノードはworld_size個の例のうち1つを保持し、他の例はスキップします
"auto" (デフォルト): データセットのシャード数がworld_sizeの約数である場合 (つまりdataset.num_shards % world_size == 0の場合) は、最適化された"shards"を使用します。それ以外の場合は"examples"を使用します。

各ノードで複数のワーカーを使用してデータをロードしたい場合は、torch.utils.data.DataLoaderと組み合わせて使用することもできます。

分散環境でイテラブルなデータセットをシャッフルする場合は、各ノードがどのシャードをスキップすべきかを把握できるように、すべてのノードで同じシャッフル済みシャードのリストが使用されるよう、IterableDataset.shuffle()内で固定のseedを設定してください。

ハイブリッドシャード型データ並列 (HSDP)

torch.distributed.device_mesh.DeviceMesh は、並列化戦略ごとに1次元ずつ持つランクのグリッドとして、トレーニングのトポロジーを記述します。データローディングにはデータ並列の次元のみが重要ですが、ハイブリッドシャーディングにより、これを見落としやすくなります。

"fsdp"（通常はノード内）: FSDPグループでは、各ランクがモデル状態（パラメータ、勾配、オプティマイザ状態）のシャードのみを保持し、モジュールの計算直前にそのモジュールの完全なパラメータを一時的にオールギャザーします。トレーニングはデータ並列のままです。各ランクが自分のバッチに対して完全なフォワードパスとバックワードパスを実行するため、この次元はDDPと同様に、ランクごとに異なるデータが必要です。
"dp_replicate"（通常はノード間）: そのシャード化されたモデルの重複コピーで、レプリカ間の追加の all-reduce によって同期が保たれます。レプリカは異なるデータを処理します。
"tp"（ノード内）: 各レイヤーの重みを分割して、1つのバッチを一緒に計算するテンソル並列ランクです。データ並列ではなく、同じ例を読み込む必要があります。

データ並列度は、シャーディング（FSDP）次元とレプリケーション次元の積であり、メッシュに"tp"のような次元があると、WORLD_SIZEより小さくなります。datasets.distributed.split_dataset_by_node()には、その積と、それらの2つの次元にわたってフラット化されたランクを渡します。

from torch.distributed.device_mesh import init_device_mesh
from datasets.distributed import split_dataset_by_node

# Hybrid sharding spreads data parallelism over two mesh dimensions, but both
# need different data per rank, so `dp_size` is their product.
mesh = init_device_mesh("cuda", (2, 2, 2), mesh_dim_names=("dp_replicate", "fsdp", "tp"))
dp_replicate_mesh, fsdp_mesh = mesh["dp_replicate"], mesh["fsdp"]

dp_size = dp_replicate_mesh.size() * fsdp_mesh.size()
dp_rank = dp_replicate_mesh.get_local_rank() * shard.size() + fsdp_mesh.get_local_rank()

ds = split_dataset_by_node(ds, rank=dp_rank, world_size=dp_size)

ここではdp_sizeは2 * 2 = 4です。各データ並列ランクはデータセットの4分の1を読み取り、各"tp"ペアの2つのランクは同じ例を読み取ります。

Update on GitHub

←ストリーム TensorFlow との併用→