Datasets documentation
Uso con JAX
Uso con JAX
Este documento es una introducción rápida al uso de datasets con JAX, con un enfoque particular en cómo obtener objetos jax.Array de nuestros datasets y cómo usarlos para entrenar modelos de JAX.
Se necesitan
jaxyjaxlibpara reproducir el código anterior, así que asegúrate de instalarlos comopip install datasets[jax].
Formato del conjunto de datos
Por defecto, los conjuntos de datos devuelven objetos de Python estándar: enteros, flotantes, cadenas, listas, etc., y los objetos de cadena y binarios no se modifican, ya que JAX solo admite números.
Para obtener arrays de JAX (similares a numpy) en su lugar, puedes establecer el formato del conjunto de datos como 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)}Un objeto Dataset es un envoltorio de una tabla Arrow, que permite lecturas rápidas desde los arrays del dataset hacia arrays de JAX.
Tenga en cuenta que el mismo procedimiento exacto se aplica a los objetos DatasetDict, de modo que
al establecer el formato de un DatasetDict como jax, todos los Dataset que lo
contienen se formatearán como 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)}Otro aspecto que debes tener en cuenta es que el formato no se aplica hasta que accedes realmente a los datos. Por lo tanto, si quieres obtener un array de JAX a partir de un dataset, primero debes acceder a los datos; de lo contrario, el formato permanecerá igual.
Por último, para cargar los datos en el dispositivo de tu elección, puedes especificar el argumento device, pero ten en cuenta que jaxlib.xla_extension.Device no es compatible, ya que no es serializable ni con pickle ni con dill, por lo que tendrás que usar su identificador de cadena en su lugar:
>>> 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]
TrueTenga en cuenta que si no se proporciona el argumento device a with_format, se utilizará el dispositivo predeterminado, que es jax.devices()[0].
Arreglos de N dimensiones
Si su conjunto de datos consiste en matrices de N dimensiones, verá que por defecto se consideran el mismo tensor si la forma es fija:
>>> 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)]}Sin embargo, esta lógica a menudo requiere comparaciones de forma y copias de datos lentas.
Para evitarlo, debes usar explícitamente el tipo de función Array y especificar la forma de tus tensores:
>>> 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)}Otros tipos de características
Los datos de ClassLabel se convierten correctamente en matrices:
>>> 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)}Los objetos de cadena y binarios no se modifican, ya que JAX solo admite números.
Los tipos de características Image y Audio también son compatibles.
Para usar el tipo de función Image, deberá instalar el extra
visioncomopip 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)}Para usar el tipo de función Audio, debes instalar el extra
audiocomopip 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)Carga de datos
JAX no tiene capacidades de carga de datos integradas, por lo que necesitará usar una biblioteca como PyTorch para cargar sus datos con un DataLoader o TensorFlow usando un tf.data.Dataset. Citando la documentación de JAX sobre este tema:
“JAX se centra de forma muy precisa en las transformaciones de programas y en NumPy respaldado por aceleradores, por lo que no incluimos la carga ni el procesamiento de datos en la biblioteca JAX. Ya existen muchos cargadores de datos excelentes, así que simplemente úselos en lugar de reinventar la rueda. Tomaremos el cargador de datos de PyTorch y crearemos un pequeño adaptador para que funcione con matrices de NumPy.”.
Por eso el formato JAX en datasets es tan útil, ya que permite utilizar cualquier modelo del HuggingFace Hub con JAX, sin tener que preocuparse por la parte de carga de datos.
Usando with_format('jax')
La forma más sencilla de obtener matrices de JAX de un conjunto de datos es utilizar el método with_format('jax'). Supongamos que queremos entrenar una red neuronal con el conjunto de datos MNIST disponible en HuggingFace Hub en 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)}Una vez que se establece el formato, podemos pasar el conjunto de datos al modelo de JAX en lotes utilizando el método Dataset.iter():
>>> for epoch in range(epochs):
... for batch in ds["train"].iter(batch_size=32):
... x, y = batch["image"], batch["label"]
... ...