Datasets documentation

JAX ile kullanma

Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

JAX ile kullanma

Bu belge, datasets kütüphanesinin JAX ile kullanımına hızlı bir giriş niteliğindedir; odak noktası, veri setlerimizden jax.Array nesnelerini nasıl elde edeceğimiz ve bunları JAX modellerini eğitmek için nasıl kullanacağımızdır.

Yukarıdaki kodu yeniden oluşturmak için jax ve jaxlib gereklidir, bu nedenle bunları pip install datasets[jax] olarak yüklediğinizden emin olun.

Veri seti biçimi

Varsayılan olarak veri kümeleri, JAX yalnızca sayıları desteklediğinden ve dizge ile ikili nesneler değiştirilmeden bırakıldığından, tamsayı, kayan noktalı sayı, dizge, liste gibi standart Python nesnelerini döndürür.

Bunun yerine JAX dizilerini (numpy benzeri) almak için veri setinin biçimini jax olarak ayarlayabilirsiniz:

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

Bir Dataset nesnesi, Arrow tablosunun bir sarmalayıcısıdır ve veri setindeki dizilerden JAX dizilerine hızlı okumalara olanak tanır.

DatasetDict nesnelerine de tam olarak aynı prosedürün uygulandığını unutmayın; bu nedenle bir DatasetDict'in biçimi jax olarak ayarlandığında, içindeki tüm Datasetler jax olarak biçimlendirilir:

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

Dikkate almanız gereken bir diğer husus, biçimlendirmenin verileri gerçekten erişene kadar uygulanmamasıdır. Dolayısıyla bir veri setinden JAX dizisi almak istiyorsanız, önce verilere erişmeniz gerekir; aksi takdirde biçim aynı kalır.

Son olarak, verileri tercih ettiğiniz cihaza yüklemek için device argümanını belirleyebilirsiniz; ancak jaxlib.xla_extension.Device, ne pickle ne de dill ile serileştirilebildiği için desteklenmez, bu nedenle bunun yerine dizge tanımlayıcısını kullanmanız gerekir:

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

device argümanı with_format işlevine sağlanmazsa, varsayılan cihaz olan jax.devices()[0] kullanılacaktır.

N boyutlu diziler

Veri setiniz N boyutlu dizilerden oluşuyorsa, şekil sabit olduğunda varsayılan olarak bunların aynı tensör olarak ele alındığını göreceksiniz:

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

Ancak bu mantık genellikle yavaş şekil karşılaştırmaları ve veri kopyalamaları gerektirir. Bunu önlemek için Array özellik türünü açıkça kullanmalı ve tensörlerinizin şeklini belirtmelisiniz:

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

Diğer özellik türleri

ClassLabel verileri düzgün şekilde dizilere dönüştürülür:

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

JAX yalnızca sayıları desteklediği için string ve ikili nesneler değiştirilmez.

Image ve Audio özellik türleri de desteklenir.

Image özellik türünü kullanmak için vision ek paketini pip install datasets[vision] olarak yüklemeniz gerekir.

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

Audio özellik türünü kullanmak için audio ek paketini pip install datasets[audio] olarak yüklemeniz gerekir.

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

Veri yükleme

JAX yerleşik veri yükleme özelliklerine sahip değildir, bu nedenle verilerinizi PyTorch ile bir DataLoader veya TensorFlow ile bir tf.data.Dataset kullanarak yüklemeniz gerekir. Bu konudaki JAX belgeleri şu ifadeleri kullanıyor: “JAX, program dönüşümlerine ve hızlandırıcı destekli NumPy’ye odaklanmıştır, bu nedenle JAX kütüphanesine veri yükleme veya işleme dahil etmiyoruz. Zaten harika veri yükleyiciler mevcut, bu yüzden her şeyi yeniden icat etmek yerine onları kullanalım. PyTorch’un veri yükleyicisini alacağız ve NumPy dizileriyle çalışması için küçük bir ara katman (shim) oluşturacağız.”

Bu nedenle datasets içindeki JAX biçimlendirme çok faydalıdır; çünkü veri yükleme kısmıyla uğraşmak zorunda kalmadan HuggingFace Hub’daki herhangi bir modeli JAX ile kullanmanızı sağlar.

with_format('jax') ile kullanma

Bir veri setinden JAX dizilerini çıkarmak için en kolay yol with_format('jax') yöntemini kullanmaktır. HuggingFace Hub'da https://huggingface.co/datasets/ylecun/mnist adresinde bulunan MNIST veri seti üzerinde bir sinir ağı eğitmek istediğimizi varsayalım.

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

Biçim ayarlandıktan sonra, veri setini Dataset.iter() yöntemini kullanarak JAX modeline partiler halinde verebiliriz:

>>> for epoch in range(epochs):
...     for batch in ds["train"].iter(batch_size=32):
...         x, y = batch["image"], batch["label"]
...         ...
Update on GitHub