Datasets documentation
Gunakan dengan JAX
Gunakan dengan JAX
Dokumen ini adalah pengantar singkat tentang penggunaan datasets dengan JAX, dengan fokus khusus pada cara mendapatkan objek jax.Array dari dataset kami, dan cara menggunakannya untuk melatih model JAX.
jaxdanjaxlibdiperlukan untuk mereproduksi kode di atas, jadi pastikan Anda menginstalnya sebagaipip install datasets[jax].
Format dataset
Secara default, dataset mengembalikan objek Python biasa: bilangan bulat, bilangan desimal, string, daftar, dll., dan objek string serta biner tidak diubah, karena JAX hanya mendukung angka.
Untuk mendapatkan array JAX (mirip numpy) sebagai gantinya, Anda dapat mengatur format dataset menjadi jax:
>>> from datasets import Dataset
>>> data = [[1, 2], [3, 4]]
>>> ds = Dataset.from_dict({"data": data})
>>> ds = ds.with_format("jax")
>>> ds[0]
{'data': DeviceArray([1, 2], dtype=int32)}
>>> ds[:2]
{'data': DeviceArray([
[1, 2],
[3, 4]], dtype=int32)}Objek Dataset adalah pembungkus dari tabel Arrow, yang memungkinkan pembacaan cepat dari array dalam dataset ke array JAX.
Perhatikan bahwa prosedur yang persis sama berlaku untuk objek DatasetDict, sehingga
ketika format DatasetDict diatur ke jax, semua Dataset di dalamnya
akan diformat sebagai jax:
>>> from datasets import DatasetDict
>>> data = {"train": {"data": [[1, 2], [3, 4]]}, "test": {"data": [[5, 6], [7, 8]]}}
>>> dds = DatasetDict.from_dict(data)
>>> dds = dds.with_format("jax")
>>> dds["train"][:2]
{'data': DeviceArray([
[1, 2],
[3, 4]], dtype=int32)}Hal lain yang perlu Anda pertimbangkan adalah bahwa format tidak diterapkan sampai Anda benar-benar mengakses data. Jadi, jika Anda ingin mendapatkan array JAX dari dataset, Anda perlu mengakses data terlebih dahulu, jika tidak format akan tetap sama.
Terakhir, untuk memuat data ke perangkat pilihan Anda, Anda dapat menentukan argumen device, namun perhatikan bahwa jaxlib.xla_extension.Device tidak didukung karena tidak dapat diserialisasi baik dengan pickle maupun dill, sehingga Anda perlu menggunakan pengenal string-nya sebagai gantinya:
>>> import jax
>>> from datasets import Dataset
>>> data = [[1, 2], [3, 4]]
>>> ds = Dataset.from_dict({"data": data})
>>> device = str(jax.devices()[0]) # Not casting to `str` before passing it to `with_format` will raise a `ValueError`
>>> ds = ds.with_format("jax", device=device)
>>> ds[0]
{'data': DeviceArray([1, 2], dtype=int32)}
>>> ds[0]["data"].device()
TFRT_CPU_0
>>> assert ds[0]["data"].device() == jax.devices()[0]
TruePerhatikan bahwa jika argumen device tidak diberikan ke with_format, maka perangkat default yang digunakan adalah jax.devices()[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("jax")
>>> ds[0]
{'data': Array([[1, 2],
[3, 4]], dtype=int32)}>>> 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("jax")
>>> ds[0]
{'data': [Array([1, 2], dtype=int32), Array([3], dtype=int32)]}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("jax")
>>> ds[0]
{'data': Array([[1, 2],
[3, 4]], dtype=int32)}
>>> ds[:2]
{'data': Array([[[1, 2],
[3, 4]],
[[5, 6],
[7, 8]]], dtype=int32)}Jenis fitur lainnya
Data ClassLabel dikonversi dengan benar menjadi array:
>>> 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("jax")
>>> ds[:3]
{'label': DeviceArray([0, 0, 1], dtype=int32)}Objek string dan biner tidak berubah, karena JAX hanya mendukung angka.
Jenis fitur Image dan Audio juga didukung.
Untuk menggunakan tipe fitur Image, Anda perlu menginstal ekstra
visionsebagaipip install datasets[vision].
>>> from datasets import Dataset, Features, Image
>>> images = ["path/to/image.png"] * 10
>>> features = Features({"image": Image()})
>>> ds = Dataset.from_dict({"image": images}, features=features)
>>> ds = ds.with_format("jax")
>>> ds[0]["image"].shape
(512, 512, 3)
>>> ds[0]
{'image': DeviceArray([[[ 255, 255, 255],
[ 255, 255, 255],
...,
[ 255, 255, 255],
[ 255, 255, 255]]], dtype=uint8)}
>>> ds[:2]["image"].shape
(2, 512, 512, 3)
>>> ds[:2]
{'image': DeviceArray([[[[ 255, 255, 255],
[ 255, 255, 255],
...,
[ 255, 255, 255],
[ 255, 255, 255]]]], dtype=uint8)}Untuk menggunakan tipe fitur Audio, Anda perlu menginstal ekstra
audiosebagaipip install datasets[audio].
>>> from datasets import Dataset, Features, Audio
>>> audio = ["path/to/audio.wav"] * 10
>>> features = Features({"audio": Audio()})
>>> ds = Dataset.from_dict({"audio": audio}, features=features)
>>> ds = ds.with_format("jax")
>>> ds[0]["audio"]["array"]
DeviceArray([-0.059021 , -0.03894043, -0.00735474, ..., 0.0133667 ,
0.01809692, 0.00268555], dtype=float32)
>>> ds[0]["audio"]["sampling_rate"]
DeviceArray(44100, dtype=int32, weak_type=True)Memuat data
JAX tidak memiliki kemampuan pemuatan data bawaan, sehingga Anda perlu menggunakan pustaka seperti PyTorch untuk memuat data Anda menggunakan DataLoader atau TensorFlow menggunakan tf.data.Dataset. Mengutip dokumentasi JAX tentang topik ini:
“JAX sangat berfokus pada transformasi program dan NumPy yang didukung akselerator, jadi kami tidak menyertakan pemuatan atau pengolahan data di pustaka JAX. Sudah ada banyak pemuat data yang bagus di luar sana, jadi mari kita gunakan saja alih-alih menciptakan ulang apa pun. Kami akan mengambil pemuat data PyTorch, dan membuat shim kecil agar dapat bekerja dengan array NumPy.”.
Itulah mengapa format JAX di datasets sangat berguna, karena memungkinkan Anda menggunakan model apa pun dari HuggingFace Hub dengan JAX, tanpa perlu khawatir tentang bagian pemuatan data.
Menggunakan with_format('jax')
Cara termudah untuk mendapatkan array JAX dari dataset adalah dengan menggunakan metode with_format('jax'). Asumsikan
bahwa kita ingin melatih jaringan saraf pada dataset MNIST yang tersedia
di HuggingFace Hub di https://huggingface.co/datasets/ylecun/mnist.
>>> from datasets import load_dataset
>>> ds = load_dataset("ylecun/mnist")
>>> ds = ds.with_format("jax")
>>> ds["train"][0]
{'image': DeviceArray([[ 0, 0, 0, ...],
[ 0, 0, 0, ...],
...,
[ 0, 0, 0, ...],
[ 0, 0, 0, ...]], dtype=uint8),
'label': DeviceArray(5, dtype=int32)}Setelah format diatur, kita dapat memberikan dataset ke model JAX dalam batch menggunakan metode Dataset.iter():
>>> for epoch in range(epochs):
... for batch in ds["train"].iter(batch_size=32):
... x, y = batch["image"], batch["label"]
... ...