Datasets documentation

Gunakan dengan JAX

You are viewing main version, which requires installation from source. If you'd like regular pip install, checkout the latest stable version (v4.8.4).
Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

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.

jax dan jaxlib diperlukan untuk mereproduksi kode di atas, jadi pastikan Anda menginstalnya sebagai pip 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]
True

Perhatikan 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 vision sebagai pip 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 audio sebagai pip 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"]
...         ...
Update on GitHub