Gunakan dengan PyTorch

Dokumen ini adalah pengantar singkat tentang penggunaan datasets dengan PyTorch, dengan fokus khusus pada cara mendapatkan objek torch.Tensor dari dataset kami, serta cara menggunakan DataLoader PyTorch dan Dataset Hugging Face dengan kinerja terbaik.

Format dataset

Secara default, dataset mengembalikan objek python biasa: bilangan bulat, bilangan desimal, string, daftar, dan sebagainya.

Untuk mendapatkan tensor PyTorch, Anda dapat mengatur format dataset menjadi pytorch menggunakan Dataset.with_format():

>>> 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]])}

Objek Dataset adalah pembungkus dari tabel Arrow, yang memungkinkan pembacaan cepat tanpa salinan (zero-copy) dari array dalam dataset ke tensor PyTorch.

Untuk memuat data sebagai tensor di GPU, tentukan argumen 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')}

Array N-dimensi

Jika dataset Anda terdiri dari array N-dimensi, Anda akan melihat bahwa secara default mereka dianggap sebagai tensor yang sama jika bentuknya tetap:

>>> 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])]}

Namun logika ini sering kali memerlukan perbandingan bentuk (shape) dan penyalinan data yang lambat. Untuk menghindarinya, Anda harus secara eksplisit menggunakan tipe fitur Array dan menentukan bentuk tensor Anda:

>>> 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]]])}

Jenis fitur lainnya

Data ClassLabel dikonversi dengan benar menjadi tensor:

>>> 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])}

Objek string dan biner tidak berubah, karena PyTorch hanya mendukung angka.

Jenis fitur Image dan Audio juga didukung.

Untuk menggunakan tipe fitur Image, Anda perlu menginstal ekstra vision sebagai 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)}

Untuk menggunakan tipe fitur Audio, Anda perlu menginstal ekstra audio sebagai 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)

Memuat data

Seperti objek torch.utils.data.Dataset, sebuah Dataset dapat diteruskan langsung ke DataLoader PyTorch:

>>> 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])}

Optimalkan pemuatan data

Ada beberapa cara untuk meningkatkan kecepatan pemuatan data Anda, yang dapat menghemat waktu, terutama jika Anda bekerja dengan dataset besar. PyTorch menawarkan pemuatan data yang diparalelkan, pengambilan batch indeks alih-alih secara individual, dan streaming untuk mengiterasi dataset tanpa mengunduhnya ke disk.

Gunakan beberapa Workers

Anda dapat memparalelkan pemuatan data dengan argumen num_workers dari DataLoader PyTorch dan mendapatkan throughput yang lebih tinggi.

Di balik layar, DataLoader memulai num_workers proses. Setiap proses memuat ulang dataset yang diteruskan ke DataLoader dan digunakan untuk menanyakan contoh. Memuat ulang dataset di dalam worker tidak akan memenuhi RAM Anda, karena hal itu hanya memetakan ulang dataset dari disk Anda ke memori.

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

Alirkan data

Alirkan dataset dengan memuatnya sebagai IterableDataset. Hal ini memungkinkan Anda untuk mengiterasi dataset jarak jauh secara bertahap tanpa mengunduhnya ke disk, atau mengiterasi file data lokal. Pelajari lebih lanjut tentang jenis dataset mana yang paling sesuai untuk kasus penggunaan Anda dalam panduan pilihan antara dataset reguler atau dataset iterable.

Dataset iterable dari datasets mewarisi dari torch.utils.data.IterableDataset sehingga Anda dapat meneruskannya ke 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)

Jika dataset dipecah menjadi beberapa shard (yaitu jika dataset terdiri dari beberapa file data), maka Anda dapat melakukan streaming secara paralel menggunakan 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)

Dalam kasus ini, setiap worker diberikan subset dari daftar shard untuk di-stream.

Checkpoint dan lanjutkan

Jika Anda memerlukan DataLoader yang dapat di-checkpoint dan dilanjutkan di tengah pelatihan, Anda dapat menggunakan StatefulDataLoader dari torchdata:

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

Hal ini dimungkinkan berkat IterableDataset.state_dict() dan IterableDataset.load_state_dict().

Distribusi

Untuk membagi dataset Anda di antara node pelatihan Anda, Anda dapat menggunakan 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"]))

Ini berlaku untuk dataset bergaya peta dan dataset yang dapat diiterasi. Dataset dipecah untuk node pada rank rank dalam kumpulan node berukuran world_size.

Untuk dataset bergaya peta:

Setiap node diberi alokasi sepotong data, misalnya rank 0 diberikan sepotong pertama dari dataset.

Untuk dataset yang dapat diiterasi:

Strategi strategy pemecahan dapat berupa "auto", "shards", atau "examples":

• "shards": shards didistribusikan secara merata di antara node
• "examples": setiap node menyimpan 1 contoh dari world_size, melewati contoh lainnya
• "auto" (default): menggunakan "shards" jika dataset memiliki jumlah shard yang merupakan faktor dari world_size (yaitu jika dataset.num_shards % world_size == 0), yang merupakan opsi paling teroptimasi. Jika tidak, menggunakan "examples".

Hal ini juga dapat digabungkan dengan torch.utils.data.DataLoader jika Anda ingin setiap node menggunakan beberapa worker untuk memuat data.

Jika Anda mengacak dataset iterable dalam setup terdistribusi, pastikan untuk mengatur seed yang tetap di IterableDataset.shuffle() sehingga daftar shard yang diacak yang sama digunakan di setiap node untuk mengetahui shard mana yang harus dilewati oleh node tersebut.

Paralel data sharded hibrida (HSDP)

torch.distributed.device_mesh.DeviceMesh menggambarkan topologi pelatihan Anda sebagai grid rank, dengan satu dimensi untuk setiap strategi paralelisasi. Hanya dimensi paralelisasi data yang relevan untuk pemuatan data, dan sharding hibrida membuatnya mudah terlewat:

• "fsdp" (biasanya intra-node): grup FSDP di mana setiap rank hanya menyimpan shard dari state model (parameter, gradien, dan state optimizer) serta secara sementara melakukan all-gather parameter penuh dari sebuah modul tepat sebelum melakukan komputasi di atasnya. Pelatihan tetap data-parallel: setiap rank menjalankan forward dan backward pass penuh pada batch-nya sendiri, sehingga dimensi ini membutuhkan data yang berbeda per rank, sama seperti DDP.
• "dp_replicate" (biasanya antar-node): salinan duplikat dari model yang di-shard tersebut, yang disinkronkan oleh all-reduce tambahan antar-replika. Replika juga memproses data yang berbeda.
• "tp" (intra-node): rank paralel tensor yang membagi bobot setiap lapisan untuk menghitung satu batch secara bersama-sama. Mereka bukan paralel data dan harus membaca contoh yang sama.

Derajat paralel data adalah hasil kali dimensi sharding (FSDP) dan replikasi, yang lebih kecil dari WORLD_SIZE begitu mesh memiliki dimensi seperti "tp". Berikan hasil kali tersebut dan rank yang diratakan di atas kedua dimensi itu ke datasets.distributed.split_dataset_by_node():

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)

Di sini dp_size adalah 2 * 2 = 4: setiap rank data-parallel membaca seperempat dari dataset, dan dua rank dari setiap pasangan "tp" membaca contoh yang sama.