Différences entre Dataset et IterableDataset

Il existe deux types d’objets de jeu de données : Dataset et IterableDataset. Le type de jeu de données que vous choisissez d’utiliser ou de créer dépend de la taille du jeu de données. En général, un IterableDataset est idéal pour les grands jeux de données (pensez à des centaines de Go !) en raison de son comportement paresseux et de ses avantages en termes de vitesse, tandis qu’un Dataset est parfait pour tout le reste. Cette page compare les différences entre un Dataset et un IterableDataset pour vous aider à choisir l’objet de jeu de données qui vous convient.

Téléchargement et streaming

Lorsque vous disposez d'un Dataset classique, vous pouvez y accéder à l'aide de my_dataset[0]. Cela permet un accès aléatoire aux lignes. Ces jeux de données sont également appelés jeux de données « map-style ». Par exemple, vous pouvez télécharger ImageNet-1k de cette manière et accéder à n'importe quelle ligne :

from datasets import load_dataset

imagenet = load_dataset("timm/imagenet-1k-wds", split="train")  # downloads the full dataset
print(imagenet[0])

Mais il y a une réserve : vous devez avoir l’ensemble du jeu de données stocké sur votre disque ou en mémoire, ce qui vous empêche d’accéder à des jeux de données plus volumineux que le disque. Comme cela peut devenir incommode pour les grands jeux de données, un autre type de jeu de données existe, le IterableDataset. Lorsque vous disposez d’un IterableDataset, vous pouvez y accéder à l’aide d’une boucle for pour charger les données progressivement au fur et à mesure que vous parcourez le jeu de données. De cette façon, seule une petite fraction d’exemples est chargée en mémoire, et vous n’écrivez rien sur le disque.

Par exemple, vous pouvez diffuser le jeu de données ImageNet-1k sans le télécharger sur le disque :

from datasets import load_dataset

imagenet = load_dataset("timm/imagenet-1k-wds", split="train", streaming=True)  # will start loading the data when iterated over
for example in imagenet:
    print(example)
    break

Le streaming permet de lire des données en ligne sans écrire aucun fichier sur le disque. Par exemple, vous pouvez streamer des jeux de données composés de plusieurs shards, chacun d'entre eux faisant des centaines de gigaoctets comme C4 ou LAION-2B. En savoir plus sur la façon de streamer un jeu de données dans le Dataset Streaming Guide.

Ce n’est pourtant pas la seule différence, car le comportement « paresseux » d’un IterableDataset est également présent lors de la création et du traitement du jeu de données.

Création de jeux de données de type map et de jeux de données itérables

Vous pouvez créer un Dataset à l'aide de listes ou de dictionnaires, et les données sont entièrement converties en Arrow afin de pouvoir accéder facilement à n'importe quelle ligne :

my_dataset = Dataset.from_dict({"col_1": [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]})
print(my_dataset[0])

En revanche, pour créer un IterableDataset, vous devez fournir une méthode « paresseuse » pour charger les données. En Python, nous utilisons généralement des fonctions générateurs. Ces fonctions yield un exemple à la fois, ce qui signifie que vous ne pouvez pas accéder à une ligne en la tranchant comme un Dataset classique :

def my_generator(n):
    for i in range(n):
        yield {"col_1": i}

my_iterable_dataset = IterableDataset.from_generator(my_generator, gen_kwargs={"n": 10})
for example in my_iterable_dataset:
    print(example)
    break

Chargement des fichiers locaux en totalité et de manière progressive

Il est possible de convertir des fichiers de données locaux ou distants en Dataset Arrow à l'aide de load_dataset() :

data_files = {"train": ["path/to/data.csv"]}
my_dataset = load_dataset("csv", data_files=data_files, split="train")
print(my_dataset[0])

Cependant, cela nécessite une étape de conversion du format CSV au format Arrow, ce qui prend du temps et de l'espace disque si votre jeu de données est volumineux.

Pour économiser de l'espace disque et éviter l'étape de conversion, vous pouvez définir un IterableDataset en diffusant directement depuis les fichiers locaux. De cette manière, les données sont lues progressivement depuis les fichiers locaux au fur et à mesure que vous itérez sur le jeu de données :

data_files = {"train": ["path/to/data.csv"]}
my_iterable_dataset = load_dataset("csv", data_files=data_files, split="train", streaming=True)
for example in my_iterable_dataset:  # this reads the CSV file progressively as you iterate over the dataset
    print(example)
    break

De nombreux formats de fichiers sont pris en charge, tels que CSV, JSONL et Parquet, ainsi que les fichiers image et audio. Vous trouverez plus d’informations dans les guides correspondants pour le chargement des jeux de données tabulaires, textuels, visuels et audio.

Traitement des données en mode immédiat et traitement des données en mode paresseux

Lorsque vous traitez un objet Dataset à l'aide de Dataset.map(), l'ensemble du jeu de données est traité immédiatement et renvoyé. Cela est similaire au fonctionnement de pandas, par exemple.

my_dataset = my_dataset.map(process_fn)  # process_fn is applied on all the examples of the dataset
print(my_dataset[0])

En revanche, en raison de la nature « paresseuse » d'un IterableDataset, l'appel de IterableDataset.map() n'applique pas votre fonction map sur l'ensemble du jeu de données. Au lieu de cela, votre fonction map est appliquée à la volée.

En conséquence, vous pouvez enchaîner plusieurs étapes de traitement et elles s’exécuteront toutes en même temps lorsque vous commencerez à itérer sur le jeu de données :

my_iterable_dataset = my_iterable_dataset.map(process_fn_1)
my_iterable_dataset = my_iterable_dataset.filter(filter_fn)
my_iterable_dataset = my_iterable_dataset.map(process_fn_2)

# process_fn_1, filter_fn and process_fn_2 are applied on-the-fly when iterating over the dataset
for example in my_iterable_dataset:  
    print(example)
    break

Mélange exact et mélange approximatif rapide

Lorsque vous mélangez un Dataset à l'aide de Dataset.shuffle(), vous appliquez un mélange exact du jeu de données. Cela fonctionne en prenant une liste d'indices [0, 1, 2, ... len(my_dataset) - 1] et en mélangeant cette liste. Ensuite, l'accès à my_dataset[0] renvoie la ligne et l'index définis par le premier élément de la correspondance d'indices qui a été mélangée :

my_dataset = my_dataset.shuffle(seed=42)
print(my_dataset[0])

Comme nous n’avons pas d’accès aléatoire aux lignes dans le cas d’un IterableDataset, nous ne pouvons pas utiliser une liste d’indices mélangés pour accéder à une ligne à une position arbitraire. Cela empêche l’utilisation d’un mélange exact. À la place, un mélange approximatif rapide est utilisé dans IterableDataset.shuffle(). Il utilise un tampon de mélange pour échantillonner des exemples aléatoires de manière itérative à partir du jeu de données. Comme le jeu de données est toujours lu de manière itérative, il offre d’excellentes performances en termes de vitesse :

my_iterable_dataset = my_iterable_dataset.shuffle(seed=42, buffer_size=100)
for example in my_iterable_dataset:
    print(example)
    break

Cependant, l’utilisation d’un tampon de mélange ne suffit pas à assurer un mélange satisfaisant pour l’entraînement de modèles d’apprentissage automatique. Ainsi, IterableDataset.shuffle() mélange également les fragments du jeu de données si celui-ci est constitué de plusieurs fichiers ou sources. De plus, il remplit le tampon en utilisant jusqu’à max_buffer_input_shards fragments à la fois (la valeur par défaut est 10), ce qui améliore considérablement la qualité du mélange :

# Stream from the internet
my_iterable_dataset = load_dataset("deepmind/code_contests", split="train", streaming=True)
my_iterable_dataset.num_shards  # 39

# Stream from local files
data_files = {"train": [f"path/to/data_{i}.csv" for i in range(1024)]}
my_iterable_dataset = load_dataset("csv", data_files=data_files, split="train", streaming=True)
my_iterable_dataset.num_shards  # 1024

# From a generator function
def my_generator(n, sources):
    for source in sources:
        for example_id_for_current_source in range(n):
            yield {"example_id": f"{source}_{example_id_for_current_source}"}

gen_kwargs = {"n": 10, "sources": [f"path/to/data_{i}" for i in range(1024)]}
my_iterable_dataset = IterableDataset.from_generator(my_generator, gen_kwargs=gen_kwargs)
my_iterable_dataset.num_shards  # 1024

Si votre jeu de données contient peu de fichiers, vous pouvez améliorer la qualité du mélange en appelant reshard() avant le mélange. Par exemple, pour les jeux de données Parquet, cela réorganise les données en utilisant des groupes de lignes au lieu d’avoir un fichier par shard, ce qui vous donne beaucoup plus de shards à mélanger.

Différences de vitesse

Les objets Dataset classiques sont basés sur Arrow, qui offre un accès aléatoire rapide aux lignes. Grâce au mappage mémoire et au fait qu’Arrow est un format en mémoire, la lecture des données depuis le disque n’effectue pas d’appels système coûteux ni de désérialisation. Il fournit un chargement des données encore plus rapide lors de l’itération à l’aide d’une boucle for en itérant sur des lots d’enregistrements Arrow contigus.

Cependant, dès que votre Dataset possède un mappage d'indices (par exemple via Dataset.shuffle()), la vitesse peut devenir 10 fois plus lente. Cela est dû à une étape supplémentaire pour obtenir l'index de la ligne à lire à l'aide du mappage d'indices, et surtout, vous ne lisez plus des blocs de données contigus. Pour restaurer la vitesse, vous devez réécrire l'ensemble du jeu de données sur votre disque à l'aide de Dataset.flatten_indices(), ce qui supprime le mappage d'indices. Cela peut toutefois prendre beaucoup de temps en fonction de la taille de votre jeu de données :

my_dataset[0]  # fast
my_dataset = my_dataset.shuffle(seed=42)
my_dataset[0]  # up to 10x slower
my_dataset = my_dataset.flatten_indices()  # rewrite the shuffled dataset on disk as contiguous chunks of data
my_dataset[0]  # fast again

Dans ce cas, nous recommandons de passer à un IterableDataset et d'exploiter sa méthode de mélange approximatif rapide IterableDataset.shuffle(). Elle ne mélange que l'ordre des shards et ajoute un tampon de mélange à votre jeu de données, ce qui maintient la vitesse de votre jeu de données optimale. Vous pouvez également remélanger le jeu de données facilement :

for example in enumerate(my_iterable_dataset):  # fast
    pass

shuffled_iterable_dataset = my_iterable_dataset.shuffle(seed=42, buffer_size=100)

for example in enumerate(shuffled_iterable_dataset):  # as fast as before
    pass

shuffled_iterable_dataset = my_iterable_dataset.shuffle(seed=1337, buffer_size=100)  # reshuffling using another seed is instantaneous

for example in enumerate(shuffled_iterable_dataset):  # still as fast as before
    pass

Si vous utilisez votre jeu de données sur plusieurs époques, la graine effective pour mélanger l’ordre des fragments dans le tampon de mélange est seed + epoch. Cela facilite le remélange d’un jeu de données entre les époques :

for epoch in range(n_epochs):
    my_iterable_dataset.set_epoch(epoch)
    for example in my_iterable_dataset:  # fast + reshuffled at each epoch using `effective_seed = seed + epoch`
        pass

Pour relancer l’itération d’un jeu de données de type map, il suffit de passer les premiers exemples :

my_dataset = my_dataset.select(range(start_index, len(dataset)))

Mais si vous utilisez un DataLoader avec un Sampler, vous devez plutôt enregistrer l'état de votre sampler (vous avez peut-être écrit un sampler personnalisé qui permet de reprendre).

En revanche, les jeux de données itérables ne permettent pas d’accéder de manière aléatoire à un index d’exemple spécifique pour reprendre à partir de celui-ci. Vous pouvez toutefois utiliser IterableDataset.state_dict() et IterableDataset.load_state_dict() pour reprendre à partir d’un point de contrôle, de la même manière que vous le feriez pour les modèles et les optimiseurs :

>>> iterable_dataset = Dataset.from_dict({"a": range(6)}).to_iterable_dataset(num_shards=3)
>>> # save in the middle of training
>>> state_dict = iterable_dataset.state_dict()
>>> # and resume later
>>> iterable_dataset.load_state_dict(state_dict)

En interne, le jeu de données itérable suit le shard en cours de lecture et l’index de l’exemple dans le shard actuel, et il stocke ces informations dans le state_dict.

Pour reprendre depuis un point de contrôle, le jeu de données ignore tous les fragments précédemment lus afin de redémarrer depuis le fragment actuel. Il lit ensuite le fragment et ignore les exemples jusqu'à atteindre l'exemple exact du point de contrôle.

La relance d’un jeu de données est donc assez rapide, car il ne relit pas les shards déjà parcourus. Toutefois, la reprise d’un jeu de données n’est généralement pas instantanée, car il faut relire depuis le début du shard en cours et ignorer les exemples jusqu’à atteindre l’emplacement du point de contrôle.

Cela peut être utilisé avec le StatefulDataLoader de torchdata, voir streaming avec un DataLoader PyTorch.

Passer du style map au style itérable

Si vous souhaitez profiter du comportement « paresseux » d’un IterableDataset ou de ses avantages en termes de vitesse, vous pouvez convertir votre Dataset de type map en IterableDataset :

my_iterable_dataset = my_dataset.to_iterable_dataset()

Si vous souhaitez mélanger votre jeu de données ou l'utiliser avec un DataLoader PyTorch, nous recommandons de générer un IterableDataset sharded :

my_iterable_dataset = my_dataset.to_iterable_dataset(num_shards=1024)
my_iterable_dataset.num_shards  # 1024