TLDR

Cet article est une introduction à TorchData: une bibliothèque pour construire de meilleures pipelines de chargement de données dans PyTorch. Son objectif est de remplacer l'approche habituelle Dataset + DataLoader. Je vais illustrer cette introduction avec un exemple concret de notebook hébergé sur GitHub: nous préparerons et chargerons des données à partir d'un véritable ensemble de données appelé Classification des images Intel en utilisant TorchData.

Qu'est-ce que le chargement de données en apprentissage automatique?

Nous pouvons le voir comme un pipeline ou une séquence d'opérations nécessaires pour transformer les données depuis leur stockage jusqu'à ce qu'elles soient prêtes à être utilisées par un modèle pour l'entraînement ou l'inférence.

Les données peuvent être stockées dans de nombreux types de stockage possibles, tels que le disque local, les compartiments cloud ou extraites depuis une API. Après avoir lu les données, un processus arbitrairement complexe les transforme en un format prêt pour le modèle, tel que des tenseurs.

Illustration d'un pipeline de chargement de données pour l'apprentissage automatique.

Le problème avec DataLoader

Le DataLoader est monolithique. Il regroupe trop de fonctionnalités dans la même classe et rend la composition et la réutilisabilité difficiles. En effet, il existe des centaines de manières de stocker un ensemble de données donné, et chacune nécessite un DataLoader hautement personnalisé: il existe de nombreux formats primitifs tels que des fichiers tensors .pt, pickle, json..., les données peuvent être regroupées par dossiers, noms de fichiers, expressions régulières, en-têtes de fichiers..., les fichiers peuvent être compressés dans différents formats, etc.

Pour prendre en charge ces cas, nous avons besoin d'un DataLoader personnalisé ou hautement configuré. Par tous les moyens, nous préférons éviter d'écrire encore et encore le même code pour gérer toutes ces opérations spécifiques dans différents contextes.

Je code encore et encore en répertoriant les fichiers, en correspondant aux modèles, ... dans tous mes DataLoaders.

Qu'est-ce que TorchData?

TorchData nous offre une solution élégante: il fournit des primitives de chargement de données à usage unique et composable. Elles sont assemblées dans des pipelines pour correspondre à des schémas de chargement de données arbitrairement complexes afin de limiter le code personnalisé au minimum. Le concept principal de la bibliothèque est le DataPipe: c'est une implémentation renommée et réorientée du Dataset pour une utilisation composée.

Il existe deux types de DataPipe.

Tout d'abord, l'IterDataPipe représente une version mise à jour de IterDataset. Ils implémentent la méthode __iter__: vous pouvez les parcourir, mais vous ne pouvez pas accéder à leurs éléments individuellement par index. Ils sont bien adaptés aux ensembles de données en continu où les lectures aléatoires sont coûteuses.

Simple exemple: nous construisons un IterDatapipe à partir d'une plage d'entiers et les regroupons en 2 lots de nombres pairs et impairs. Pour ce faire, nous utilisons le DataPipe groupby pour les regrouper selon le résultat du modulo 2. Ensuite, nous itérons et affichons le résultat.

# IterDataPipe des 10 premiers entiers regroupés en 2 lots: pairs et impairs
pipe = (
    # Envelopper la plage dans un wrapper IterableWrapper
    dp.iter.IterableWrapper(range(10))
    # Regrouper les nombres pairs et impairs
    .groupby(lambda x: x % 2)
)
# Nous pouvons itérer sur les éléments
print("Itération complète:", list(pipe))
# Mais nous ne pouvons pas y accéder individuellement par index
# pipe[0] lèverait une exception

# Sorties:
# >> Itération complète: [[0, 2, 4, 6, 8], [1, 3, 5, 7, 9]]

Ensuite, le MapDataPipe représente une version mise à jour du MapDataset. Ils sont bien adaptés aux ensembles de données clé-valeur, où les lectures aléatoires sont bon marché. Vous pouvez les parcourir et accéder à leurs éléments individuellement par index.

Simple exemple: nous construisons un MapDataPipe à partir d'une plage d'entiers et multiplions chaque élément par 2 en utilisant un MapDataPipe fonctionnel, puis nous mélangeons la sortie. Ensuite, nous itérons, affichons le résultat et montrons que nous pouvons accéder aux éléments par index.

# MapDataPipe des 10 premiers entiers multipliés par 2 et mélangés
pipe = (
    # Envelopper la plage dans un MapDataPipe
    dp.map.SequenceWrapper(range(10))
    # Multiplier chaque nombre par 2
    .map(lambda x: x * 2)
    # Mélanger
    .shuffle()
)

# Nous pouvons itérer sur les valeurs
print("Itération complète:", list(pipe))
# Nous pouvons également accéder aux éléments individuellement en fonction de leur index
print("Accès basé sur l'index:", pipe[0], pipe[9])

# Sorties
# >> Itération complète: [4, 8, 18, 10, 14, 12, 16, 6, 0, 2]
# >> Accès basé sur l'index: 4 2

Exemple concret: Classification des images Intel

Voyons maintenant un exemple concret: le chargement de l'ensemble de données Classification des images Intel d'abord avec le DataLoader habituel, puis en utilisant la bibliothèque TorchData. Vous pouvez lire l'intégralité du code depuis le dépôt.

Exemples d'ensemble de données de Classification des images Intel.

Nous commençons par installer les outils dont nous avons besoin pour ce projet. Beaucoup sont déjà présents dans la configuration de Google Colab, mais nous devons ajouter kaggle, torchvision et torchdata.

!pip install kaggle torchdata torchvision

Puisque nous utilisons un ensemble de données Kaggle, nous téléchargeons notre clé d'API générée sur le site web de Kaggle. Elle est nécessaire pour utiliser leur client Python et télécharger facilement l'ensemble de données.

from google.colab import files
files.upload()

!mkdir -p ~/.kaggle
!cp kaggle.json ~/.kaggle/
!chmod 600 ~/.kaggle/kaggle.json

Ensuite, nous obtenons l'ensemble de données avec une simple commande et le décompressons dans un dossier nommé data.

!kaggle datasets download puneet6060/intel-image-classification

Si tout s'est bien passé, vous devriez voir 6 dossiers, un pour chaque classe, dans les dossiers train et test sous data.

!ls data/seg_test/seg_test

# Sortie
# >> buildings forest glacier mountain sea street

Configuration

Nous commençons par importer certaines bibliothèques que nous utiliserons plus tard pour le projet: glob, itertools, pathlib, torch et torchvision.

import glob
import itertools as it
from pathlib import Path

import torch
import torchvision

Ensuite, nous créons certaines utilitaires pour plus tard dans le projet.

Nous avons 2 divisions possibles, train et test, et nous créons un dictionnaire pour convertir une division en son chemin correspondant sur le disque.

En ce qui concerne l'ensemble de données, nous avons 6 classes différentes et nous créons un dictionnaire pour convertir chaque classe en un label entier.

Toutes les images n'ont pas les mêmes dimensions, nous avons donc besoin d'une transformation pour les redimensionner à une taille fixe. 150 par 150. Enfin, nous écrivons une fonction pour nous donner une classe d'image en fonction du chemin.

Le premier dossier parent est le nom de la classe, nous prenons donc simplement la racine du premier parent. La racine est simplement la dernière partie du chemin.

# Convertir le nom de la division en chemin de dossier
split_to_path = {
    "train": "data/seg_train/seg_train",
    "test": "data/seg_test/seg_test"
}

# Convertir le nom de la classe en label entier
name_to_label = {
    "buildings": 0,
    "forest": 1,
    "glacier": 2,
    "mountain": 3,
    "sea": 4,
    "street": 5,
}


# Transformations d'image pour obtenir toutes les images de taille 150 x 150
transforms = torch.nn.Sequential(
    torchvision.transforms.Resize((150, 150))
)

def img_path_to_label(path: str):
    """Fonction pour obtenir la classe à partir du chemin du fichier"""
    name = Path(path).parents[0].stem
    return name_to_label[name]

Chargement de données habituel: Dataset et DataLoader

Nous avons tout ce dont nous avons besoin pour commencer. Nous commençons par l'implémentation habituelle du chargement de données avec un Dataset et un DataLoader.

Nous importons d'abord les classes Dataset et DataLoader de torch.utils.data.

from torch.utils.data import Dataset, DataLoader

La première étape consiste à créer notre IntelDataset. Dans la méthode __init__, nous récupérons simplement le chemin en fonction de la division en utilisant la fonction définie précédemment.

Nous définissons une méthode appelée list files, qui liste simplement toutes les images au chemin donné.

L'interface Dataset nécessite de mettre en œuvre la méthode __len__ qui renvoie la taille. Ici, il s'agit simplement du nombre d'images.

La dernière étape consiste à implémenter getitem qui récupère un tuple d'image et de label en fonction d'un index: nous obtenons le chemin du fichier à l'index reçu, puis nous chargeons l'image en utilisant torchvision, nous obtenons le label de l'image en utilisant notre fonction d'utilité, et enfin nous obtenons l'image redimensionnée avec le label correspondant.

class IntelDataset(Dataset):
    """Classe pour représenter la Classification des images Intel en tant que Dataset"""
    def __init__(self, split: str):
        # Obtenir le chemin de division (train ou test) à partir du nom de la division.
        self.path = split_to_path[split]

    def _list_files(self):
        """Lister toutes les images"""
        return list(glob.glob(f"{self.path}/**/*.jpg"))

    def __len__(self):
        """Obtenir la longueur de l'ensemble de données"""
        return len(self._list_files())

    def __getitem__(self, idx: int):
        """Méthode pour accéder à un tuple (entrée, label) par index"""
        # Obtenir tous les chemins de fichiers
        files = self._list_files()
        # Obtenir le chemin du fichier à l'index reçu
        file_path = files[idx]
        # Charger l'image
        image = torchvision.io.read_image(file_path)
        # Obtenir le label à partir du chemin de l'image
        label = img_path_to_label(file_path)
        # Renvoyer l'image transformée avec son label
        return transforms(image), label

Nous avons notre implémentation du Dataset, nous pouvons maintenant créer le DataLoader avec mélange et une taille de lot de 10 éléments. Pour vérifier que tout fonctionne comme prévu, nous itérons sur les 5 premiers lots et affichons la taille ainsi que les labels. Nous pouvons voir ici que nous obtenons la taille du lot attendue et un ensemble de données mélangé.

# Créer le Dataset pour la division train
ds = IntelDataset("train")
# Créer le DataLoader avec mélange et mise en lots
dl = DataLoader(ds, batch_size=10, shuffle=True)

# Itérer sur les 5 premiers lots
for X, y in it.islice(dl, 5):
    print(f"Longueur du lot X: {len(X)}, longueur du lot y: {len(y)}, labels: {y}")

# Sorties:
# >> Longueur du lot X: 10, longueur du lot y: 10, labels: tensor([1, 1, 3, 1, 2, 3, 5, 4, 4, 2])
# >> Longueur du lot X: 10, longueur du lot y: 10, labels: tensor([4, 4, 4, 2, 4, 4, 5, 5, 2, 1])
# >> Longueur du lot X: 10, longueur du lot y: 10, labels: tensor([4, 3, 4, 3, 3, 2, 4, 1, 1, 4])
# >> Longueur du lot X: 10, longueur du lot y: 10, labels: tensor([1, 4, 0, 4, 4, 5, 1, 3, 5, 0])
# >> Longueur du lot X: 10, longueur du lot y: 10, labels: tensor([3, 0, 5, 5, 1, 1, 2, 2, 1, 0])

Chargement de données comme un pro: TorchData

En analysant le code ci-dessus, nous pouvons identifier certaines primitives qui devraient idéalement être implémentées dans une bibliothèque: lister les fichiers dans un dossier, mapper une fonction sur les sorties de l'ensemble de données. Avec l'approche DataLoader, toute la logique se trouve dans la méthode __getitem__.

Jetons maintenant un coup d'œil à la solution élégante proposée par TorchData: combiner des primitives modulaires avec une API fonctionnelle.

Nous commençons par importer le module datapipes.

import torchdata.datapipes as dp
from torch.utils.data import default_collate

Nous créons une fonction appelée build_datapipes qui renvoie un MapDataPipes en fonction d'un nom de division. Nous commençons par utiliser notre fonction utilitaire pour récupérer le chemin du dossier de division.

Nous commençons le pipeline en répertoriant de manière récursive tous les fichiers au chemin donné. Cela nous donne un IterDataPipe.

Ensuite, nous mappions une fonction pour renvoyer des tuples de chemin d'image et le label correspondant.

Nous préférons un MapDataPipe pour cet ensemble de données pour permettre un accès basé sur l'index. Pour cela, nous avons besoin d'un index pour chaque élément. Nous l'obtenons en utilisant le pipe enumerate qui énumère de manière paresseuse tous les éléments dans l'ordre. Nous pouvons ensuite appeler to_map_datapipe pour convertir notre IterDataPipe en MapDataPipe.

Nous pouvons maintenant lire l'image dans un tenseur avec torchvision. N'oubliez pas que nous devons redimensionner les images à une dimension fixe.

À ce stade, nous mélangeons les données en enchaînant un shuffler data pipe.

Après la mise en lots de 10 éléments, nous obtenons un pipeline qui renvoie une liste de 10 tuples. Chaque tuple contenant une image avec un label. Mais idéalement, nous voulons uniquement des tenseurs, l'un contenant le lot d'images et l'autre contenant le lot de labels. Pour ce faire, nous appliquons la fonction de regroupement par défaut de torch.utils.data. Elle transforme une liste de tuples de tenseurs en un tuple de tenseurs.

def build_datapipes(split: str):
    """Fonction pour renvoyer le DataPipe en fonction du nom de la division"""

    # Obtenir le chemin de division (train ou test) à partir du nom de division.
    path = split_to_path[split]

    return (
        # Itérer sur tous les chemins de fichiers
        dp.iter.FileLister(path, recursive=True)
        # Transformer le chemin en tuples (chemin, label)
        .map(lambda x: (x, img_path_to_label(x)))
        # Nous avons besoin d'une clé pour transformer nos IterDataPipes en MapDataPipes
        # Enumerate renverra : (index, (chemin, label)) 
        .enumerate()
        # Obtenir un MapDataPipes, c'est comme un dictionnaire avec un accès basé sur la clé
        .to_map_datapipe()
        # Lire l'image et renvoyer (tenseur d'image, label)
        .map(lambda x: (torchvision.io.read_image(x[0]), x[1]))
        # Redimensionner l'image en utilisant notre transformation (image transformée, label)
        .map(lambda x: (transforms(x[0]), x[1]))
        # Mélanger les DataPipes
        .shuffle()
        # Obtenir des lots de 10
        .batch(10)
        # Rassembler les lots. Transforme [(image, label)] en (images, labels)
        .map(lambda x: default_collate(x))
)

Et voilà, notre data pipe est prêt à charger notre ensemble de données. Nous itérons sur les 5 premiers lots pour vérifier que cela fonctionne. Et comme prévu, nous obtenons des lots de 10 éléments mélangés.

pipe = build_datapipes("train")
# Itérer sur les 5 premiers lots
for X, y in it.islice(pipe, 5):
    print(f"Longueur du lot X: {len(X)}, longueur du lot y: {len(y)}, labels: {y}")

# Sorties:
# >> Longueur du lot X: 10, longueur du lot y: 10, labels: tensor([5, 4, 0, 3, 5, 2, 1, 3, 1, 5])
# >> Longueur du lot X: 10, longueur du lot y: 10, labels: tensor([3, 4, 3, 0, 5, 0, 0, 0, 3, 5])
# >> Longueur du lot X: 10, longueur du lot y: 10, labels: tensor([4, 4, 1, 4, 2, 5, 2, 1, 3, 5])
# >> Longueur du lot X: 10, longueur du lot y: 10, labels: tensor([3, 5, 2, 0, 4, 1, 3, 0, 5, 3])
# >> Longueur du lot X: 10, longueur du lot y: 10, labels: tensor([3, 5, 4, 3, 5, 1, 2, 4, 5, 4])

La bibliothèque TorchData permet la réutilisabilité tout en préservant la flexibilité, car vous pouvez assembler des primitives pour former des pipelines complexes.

Conclusion

En quelques mots, le chargement de données est une partie importante de tout projet d'apprentissage automatique. Cependant, l'approche habituelle de PyTorch avec Dataset et DataLoader est défectueuse: le DataLoader est monolithique et rend la réutilisabilité et la composition difficiles.

TorchData est une excellente solution qui apporte des primitives de chargement de données modulaires et composables: nous écrivons moins de code personnalisé, car la bibliothèque abstrait de nombreux cas d'utilisation courants dans des DataPipes.

Essayez TorchData dans votre prochain projet et vous ne reviendrez jamais en arrière.