與 PyTorch 結合使用

本文件簡要介紹如何將 datasets 與 PyTorch 結合使用，特別關注如何從我們的資料集中獲取 torch.Tensor 物件，以及如何以最佳效能使用 PyTorch 的 DataLoader 和 Hugging Face 的 Dataset。

資料集格式

預設情況下，資料集返回常規的 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維陣列組成，您會發現如果形狀固定，它們預設被視為相同的張量

>>> 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 的資料集，並用於查詢示例。在工作程序中重新載入資料集不會佔用您的記憶體，因為它只是從您的磁碟上再次記憶體對映資料集。

>>> 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()。

分散式

要將您的資料集分佈到多個訓練節點上，您可以使用 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）資料集和可迭代（iterable）資料集。資料集會為大小為 world_size 的節點池中排名為 rank 的節點進行分割。

對於對映式資料集

每個節點都會被分配一個數據塊，例如，排名為 0 的節點會被分配資料集的第一個資料塊。

對於可迭代資料集

如果資料集的分片數量是 world_size 的因子（即，如果 dataset.num_shards % world_size == 0），那麼分片會在節點之間均勻分配，這是最最佳化的。否則，每個節點會保留 world_size 個示例中的 1 個，跳過其他示例。

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

< > 在 GitHub 上更新