Utiliser avec PyTorch

Ce document est une introduction rapide à l’utilisation de datasets avec PyTorch, avec un accent particulier sur la façon d’extraire des objets torch.Tensor de nos jeux de données, et sur la façon d’utiliser un DataLoader PyTorch et un Dataset Hugging Face pour obtenir les meilleures performances.

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.

Pour obtenir des tenseurs PyTorch à la place, vous pouvez définir le format du jeu de données sur pytorch à l'aide de Dataset.with_format() :

>>> from datasets import Dataset
>>> data = [[1, 2],[3, 4]]
>>> ds = Dataset.from_dict({"data": data})
>>> ds = ds.with_format("torch")
>>> ds[0]
{'data': tensor([1, 2])}
>>> ds[:2]
{'data': tensor([[1, 2],
         [3, 4]])}

Un objet Dataset est un conteneur pour un tableau Arrow, qui permet des lectures rapides sans copie depuis les tableaux du jeu de données vers des tenseurs PyTorch.

Pour charger les données sous forme de tenseurs sur un GPU, spécifiez l'argument device :

>>> import torch
>>> device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
>>> ds = ds.with_format("torch", device=device)
>>> ds[0]
{'data': tensor([1, 2], device='cuda:0')}

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("torch")
>>> ds[0]
{'data': tensor([[1, 2],
         [3, 4]])}

>>> 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("torch")
>>> ds[0]
{'data': [tensor([1, 2]), tensor([3])]}

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("torch")
>>> ds[0]
{'data': tensor([[1, 2],
         [3, 4]])}
>>> ds[:2]
{'data': tensor([[[1, 2],
          [3, 4]],
 
         [[5, 6],
          [7, 8]]])}

Autres types de fonctionnalités

Les données ClassLabel sont correctement converties en tenseurs :

>>> 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("torch")  
>>> ds[:3]
{'label': tensor([0, 0, 1])}

Les objets de type chaîne et binaire ne sont pas modifiés, car PyTorch 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, Audio, Image
>>> images = ["path/to/image.png"] * 10
>>> features = Features({"image": Image()})
>>> ds = Dataset.from_dict({"image": images}, features=features) 
>>> ds = ds.with_format("torch")
>>> ds[0]["image"].shape
torch.Size([512, 512, 4])
>>> ds[0]
{'image': tensor([[[255, 215, 106, 255],
         [255, 215, 106, 255],
         ...,
         [255, 255, 255, 255],
         [255, 255, 255, 255]]], dtype=torch.uint8)}
>>> ds[:2]["image"].shape
torch.Size([2, 512, 512, 4])
>>> ds[:2]
{'image': tensor([[[[255, 215, 106, 255],
          [255, 215, 106, 255],
          ...,
          [255, 255, 255, 255],
          [255, 255, 255, 255]]]], dtype=torch.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, Image
>>> audio = ["path/to/audio.wav"] * 10
>>> features = Features({"audio": Audio()})
>>> ds = Dataset.from_dict({"audio": audio}, features=features) 
>>> ds = ds.with_format("torch")  
>>> ds[0]["audio"]["array"]
tensor([ 6.1035e-05,  1.5259e-05,  1.6785e-04,  ..., -1.5259e-05,
        -1.5259e-05,  1.5259e-05])
>>> ds[0]["audio"]["sampling_rate"]
tensor(44100)

Chargement des données

Comme les objets torch.utils.data.Dataset, un Dataset peut être transmis directement à un DataLoader de PyTorch :

>>> import numpy as np
>>> from datasets import Dataset 
>>> from torch.utils.data import DataLoader
>>> data = np.random.rand(16)
>>> label = np.random.randint(0, 2, size=16)
>>> ds = Dataset.from_dict({"data": data, "label": label}).with_format("torch")
>>> dataloader = DataLoader(ds, batch_size=4)
>>> for batch in dataloader:
...     print(batch)                                                                                            
{'data': tensor([0.0047, 0.4979, 0.6726, 0.8105]), 'label': tensor([0, 1, 0, 1])}
{'data': tensor([0.4832, 0.2723, 0.4259, 0.2224]), 'label': tensor([0, 0, 0, 0])}
{'data': tensor([0.5837, 0.3444, 0.4658, 0.6417]), 'label': tensor([0, 1, 0, 0])}
{'data': tensor([0.7022, 0.1225, 0.7228, 0.8259]), 'label': tensor([1, 1, 1, 1])}

Optimiser le chargement des données

Plusieurs méthodes permettent d’accélérer le chargement de vos données, ce qui peut vous faire gagner du temps, en particulier si vous travaillez avec de grands jeux de données. PyTorch propose un chargement de données parallélisé, qui récupère des lots d’indices plutôt que de les traiter individuellement, ainsi qu’un streaming pour parcourir le jeu de données sans le télécharger sur le disque.

Utiliser plusieurs Workers

Vous pouvez paralléliser le chargement des données à l'aide de l'argument num_workers d'un DataLoader PyTorch et ainsi obtenir un débit supérieur.

Sous le capot, le DataLoader démarre num_workers processus. Chaque processus recharge le jeu de données transmis au DataLoader et est utilisé pour interroger les exemples. Recharger le jeu de données dans un travailleur n’occupe pas votre RAM, car il se contente de réassocier le jeu de données à la mémoire depuis votre disque.

>>> import numpy as np
>>> from datasets import Dataset, load_from_disk
>>> from torch.utils.data import DataLoader
>>> data = np.random.rand(10_000)
>>> Dataset.from_dict({"data": data}).save_to_disk("my_dataset")
>>> ds = load_from_disk("my_dataset").with_format("torch")
>>> dataloader = DataLoader(ds, batch_size=32, num_workers=4)

Diffuser des données

Diffusez un jeu de données en le chargeant en tant que IterableDataset. Cela vous permet d’itérer progressivement sur un jeu de données distant sans le télécharger sur le disque, ou sur des fichiers de données locaux. En savoir plus sur le type de jeu de données le plus adapté à votre cas d’usage dans le guide choisir entre un jeu de données classique et un jeu de données itérable.

Un jeu de données itérable issu de datasets hérite de torch.utils.data.IterableDataset, ce qui vous permet de le passer à un torch.utils.data.DataLoader :

>>> import numpy as np
>>> from datasets import Dataset, load_dataset
>>> from torch.utils.data import DataLoader
>>> data = np.random.rand(10_000)
>>> Dataset.from_dict({"data": data}).push_to_hub("<username>/my_dataset")  # Upload to the Hugging Face Hub
>>> my_iterable_dataset = load_dataset("<username>/my_dataset", streaming=True, split="train")
>>> dataloader = DataLoader(my_iterable_dataset, batch_size=32)

Si le jeu de données est réparti en plusieurs shards (c’est-à-dire s’il se compose de plusieurs fichiers de données), vous pouvez effectuer le streaming en parallèle à l’aide de num_workers :

>>> my_iterable_dataset = load_dataset("deepmind/code_contests", streaming=True, split="train")
>>> my_iterable_dataset.num_shards
39
>>> dataloader = DataLoader(my_iterable_dataset, batch_size=32, num_workers=4)

Dans ce cas, chaque worker reçoit un sous-ensemble de la liste des shards à streamer.

Point de contrôle et reprise

Si vous avez besoin d'un DataLoader que vous pouvez sauvegarder et reprendre en cours d'entraînement, vous pouvez utiliser le StatefulDataLoader de torchdata :

>>> from torchdata.stateful_dataloader import StatefulDataLoader
>>> my_iterable_dataset = load_dataset("deepmind/code_contests", streaming=True, split="train")
>>> dataloader = StatefulDataLoader(my_iterable_dataset, batch_size=32, num_workers=4)
>>> # save in the middle of training
>>> state_dict = dataloader.state_dict()
>>> # and resume later
>>> dataloader.load_state_dict(state_dict)

Cela est possible grâce à IterableDataset.state_dict() et IterableDataset.load_state_dict().

Distribué

Pour répartir votre jeu de données sur vos nœuds d'entraînement, vous pouvez utiliser datasets.distributed.split_dataset_by_node() :

import os
from datasets.distributed import split_dataset_by_node

ds = split_dataset_by_node(ds, rank=int(os.environ["RANK"]), world_size=int(os.environ["WORLD_SIZE"]))

Cela fonctionne à la fois pour les jeux de données de type map et les jeux de données itérables. Le jeu de données est partitionné pour le nœud de rang rank dans un pool de nœuds de taille world_size.

Pour les jeux de données de type carte :

Chaque nœud se voit attribuer un fragment de données, par exemple le rang 0 reçoit le premier fragment du jeu de données.

Pour les jeux de données itérables :

La strategy de découpage peut être "auto", "shards" ou "examples" :

• "shards" : les shards sont répartis uniformément entre les nœuds
• "examples" : chaque nœud conserve 1 exemple sur world_size et ignore les autres exemples
• "auto" (par défaut) : utilise "shards" si le nombre de shards du jeu de données est un diviseur de world_size (c’est-à-dire si dataset.num_shards % world_size == 0), ce qui est l’option la plus optimisée. Sinon, utilise "examples".

Cela peut également être combiné avec un torch.utils.data.DataLoader si vous souhaitez que chaque nœud utilise plusieurs workers pour charger les données.

Si vous mélangez votre jeu de données itérable dans un environnement distribué, assurez-vous de définir une seed fixe dans IterableDataset.shuffle() afin que la même liste mélangée de shards soit utilisée sur chaque nœud pour savoir quels shards le nœud doit ignorer.

Parallélisme de données hybride et fragmenté (HSDP)

Un torch.distributed.device_mesh.DeviceMesh décrit votre topologie d’entraînement sous forme d’une grille de rangs, avec une dimension par stratégie de parallélisme. Seules les dimensions de parallélisme de données sont pertinentes pour le chargement des données, et le fractionnement hybride les rend faciles à manquer :

• "fsdp" (généralement intra-nœud) : le groupe FSDP où chaque rang ne conserve qu'un fragment des états du modèle (paramètres, gradients et états de l'optimiseur) et effectue temporairement un all-gather des paramètres complets d'un module juste avant de calculer dessus. L'entraînement reste parallèle sur les données : chaque rang exécute un passage avant et un passage arrière complets sur son propre lot, de sorte que cette dimension nécessite des données différentes par rang, tout comme avec DDP.
• "dp_replicate" (généralement inter-nœuds) : les copies dupliquées de ce modèle shardé, maintenues synchronisées par un all-reduce supplémentaire entre les répliques. Les répliques traitent également des données différentes.
• "tp" (intra-nœud) : les rangs de parallélisme de tenseurs qui divisent les poids de chaque couche pour calculer un même lot ensemble. Ils ne sont pas en parallélisme de données et doivent lire les mêmes exemples.

Le degré de parallélisme des données est le produit des dimensions de sharding (FSDP) et de réplication, ce qui est inférieur à WORLD_SIZE dès que le maillage possède une dimension comme "tp". Fournissez à datasets.distributed.split_dataset_by_node() ce produit ainsi que le rank aplati sur ces deux dimensions :

from torch.distributed.device_mesh import init_device_mesh
from datasets.distributed import split_dataset_by_node

# Hybrid sharding spreads data parallelism over two mesh dimensions, but both
# need different data per rank, so `dp_size` is their product.
mesh = init_device_mesh("cuda", (2, 2, 2), mesh_dim_names=("dp_replicate", "fsdp", "tp"))
dp_replicate_mesh, fsdp_mesh = mesh["dp_replicate"], mesh["fsdp"]

dp_size = dp_replicate_mesh.size() * fsdp_mesh.size()
dp_rank = dp_replicate_mesh.get_local_rank() * shard.size() + fsdp_mesh.get_local_rank()

ds = split_dataset_by_node(ds, rank=dp_rank, world_size=dp_size)

Ici, dp_size vaut 2 * 2 = 4 : chaque rang parallèle de données lit un quart du jeu de données, et les deux rangs de chaque paire "tp" lisent les mêmes exemples.