與 PyTorch 搭配使用

本文件是將 datasets 與 PyTorch 搭配使用的快速入門，特別著重於如何從我們的資料集中獲取 torch.Tensor 物件，以及如何以最佳效能使用 PyTorch 的 DataLoader 和 Hugging Face 的 Dataset。

資料集格式

預設情況下，datasets 會回傳一般的 Python 物件：整數、浮點數、字串、列表等。

若要獲取 PyTorch 張量（tensor），您可以使用 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 表格的封裝，它允許從資料集中的陣列快速進行零拷貝（zero-copy）讀取，轉換為 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 維陣列組成，您會發現若形狀固定，它們預設會被視為同一個張量

>>> 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 提供了平行化資料載入、批次檢索索引（而非個別檢索），以及透過串流（streaming）來遍歷資料集，而無需將其下載至磁碟。

使用多個工作程序 (Workers)

您可以透過 PyTorch DataLoader 的 num_workers 引數來平行化資料載入，並獲得更高的吞吐量。

在底層，DataLoader 會啟動 num_workers 個程序。每個程序會重新載入傳遞給 DataLoader 的資料集，並用於查詢範例。在工作程序內重新載入資料集不會佔滿您的記憶體（RAM），因為它只是再次將資料集從磁碟進行記憶體映射（memory-mapping）。

>>> 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 來進行串流。這讓您可以逐步遍歷遠端資料集或本機資料檔案，而無需將其下載至磁碟。若要深入了解哪種資料集類型最適合您的使用案例，請參閱選擇一般資料集或可迭代資料集 (iterable dataset) 指南。

來自 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)

如果資料集被分割成多個分片 (shards)（即資料集由多個資料檔案組成），那麼您可以使用 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()。

分散式

若要將資料集分散到您的訓練節點（training nodes），可以使用 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-style) 資料集與可迭代資料集。資料集會針對 rank 為 rank 的節點，在大小為 world_size 的節點池中進行分割。

針對映射式資料集：

每個節點會被分配到一塊資料，例如 rank 0 會獲得資料集的第一塊。

針對可迭代資料集：

如果資料集的分片數量是 world_size 的倍數（即 dataset.num_shards % world_size == 0），則分片會均勻分配到各個節點，這是最優化的方式。否則，每個節點會保留 1/world_size 的範例，並跳過其他範例。

如果您希望每個節點使用多個工作程序來載入資料，這也可以與 torch.utils.data.DataLoader 結合使用。

如果您在分散式設定中對可迭代資料集進行洗牌（shuffle），請務必在 IterableDataset.shuffle() 中設定固定的 seed，這樣每個節點就會使用相同的洗牌分片列表，以確認該節點應該跳過哪些分片。

在 GitHub 上更新

資料集 (Datasets)

與 PyTorch 搭配使用

資料集格式

N 維陣列

其他特徵類型

資料載入

優化資料載入

使用多個工作程序 (Workers)

串流資料

檢查點與恢復

分散式