Utiliser avec JAX

Ce document est une introduction rapide à l'utilisation de datasets avec JAX, en mettant particulièrement l'accent sur la manière d'extraire des objets jax.Array de nos jeux de données et de les utiliser pour entraîner des modèles JAX.

jax et jaxlib sont nécessaires pour reproduire le code ci-dessus, veuillez donc vous assurer de les installer via pip install datasets[jax].

Format du jeu de données

Par défaut, les jeux de données renvoient des objets Python classiques : des entiers, des flottants, des chaînes de caractères, des listes, etc., et les objets de type chaîne et binaire sont laissés inchangés, car JAX ne prend en charge que les nombres.

Pour obtenir des tableaux JAX (de type numpy) à la place, vous pouvez définir le format du jeu de données sur 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 objet Dataset est un conteneur pour un tableau Arrow, qui permet des lectures rapides depuis les tableaux du jeu de données vers des tableaux JAX.

Notez que la procédure exacte s'applique également aux objets DatasetDict, de sorte que lorsque le format d'un DatasetDict est défini sur jax, tous les Dataset qu'il contient seront formatés en 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)}

Un autre point à prendre en compte est que le formatage n’est appliqué que lorsque vous accédez réellement aux données. Ainsi, si vous souhaitez obtenir un tableau JAX à partir d’un jeu de données, vous devez d’abord accéder aux données, sinon le format restera inchangé.

Enfin, pour charger les données sur l’appareil de votre choix, vous pouvez spécifier l’argument device, mais notez que jaxlib.xla_extension.Device n’est pas pris en charge car il n’est pas sérialisable ni avec pickle ni avec dill ; vous devrez donc utiliser son identifiant sous forme de chaîne à la place :

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

Notez que si l'argument device n'est pas fourni à with_format, l'appareil par défaut, soit jax.devices()[0], sera utilisé.

Tableaux N-dimensionnels

Si votre jeu de données est constitué de tableaux N-dimensionnels, vous constaterez qu’ils sont considérés par défaut comme le même tenseur si la forme est fixe :

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

Cependant, cette logique nécessite souvent des comparaisons de formes et des copies de données lentes. Pour l’éviter, vous devez utiliser explicitement le type de fonctionnalité Array et spécifier la forme de vos tenseurs :

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

Autres types de fonctionnalités

Les données ClassLabel sont correctement converties en tableaux :

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

Les objets de type chaîne et binaire ne sont pas modifiés, car JAX ne prend en charge que les nombres.

Les types de fonctionnalités Image et Audio sont également pris en charge.

Pour utiliser le type de fonctionnalité Image, vous devez installer l'extra vision via 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)}

Pour utiliser le type de fonctionnalité Audio, vous devez installer l'extra audio via 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)

Chargement des données

JAX ne dispose d’aucune fonctionnalité intégrée de chargement de données, il est donc nécessaire d’utiliser une bibliothèque telle que PyTorch pour charger vos données à l’aide d’un DataLoader ou TensorFlow en utilisant un tf.data.Dataset. En citant la documentation JAX sur ce sujet : « JAX est axé sur les transformations de programmes et sur NumPy adossé à des accélérateurs, nous n’incluons donc pas le chargement ou le traitement des données dans la bibliothèque JAX. Il existe déjà de nombreux chargeurs de données excellents, utilisons-les donc plutôt que de réinventer quoi que ce soit. Nous utiliserons le chargeur de données de PyTorch et créerons un petit adaptateur pour le faire fonctionner avec des tableaux NumPy. ».

C’est la raison pour laquelle le formatage JAX dans datasets est si utile, car il permet d’utiliser n’importe quel modèle du HuggingFace Hub avec JAX, sans avoir à se soucier de la partie chargement des données.

Utilisation de with_format('jax')

La façon la plus simple d’extraire des tableaux JAX d’un jeu de données est d’utiliser la méthode with_format('jax'). Supposons que nous voulons entraîner un réseau de neurones sur le jeu de données MNIST disponible sur le HuggingFace Hub à l’adresse 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)}

Une fois le format défini, nous pouvons alimenter le modèle JAX avec le jeu de données par lots à l'aide de la méthode Dataset.iter() :

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