Skip to main content

Chargement de bases de données pour l'apprentissage distribué en PyTorch

Dans cette page, nous mettons en pratique la gestion des Datasets et DataLoaders pour l'apprentissage distribué en PyTorch. Nous nous intéressons aux problématiques présentées dans la page principale sur le chargement des données.

Nous présentons ici l'usage :

La documentation se conclut sur la présentation d'un exemple complet de chargement optimisé de données, et sur une mise en pratique sur Jean Zay via un Jupyter Notebook.

remarque préliminaire

Dans cette documentation, nous ne parlerons pas des objets de type IterableDataset qui permettent de traiter des bases de données dont la structure est inconnue. Ce genre d'objet est parcouru à l'aide d'un itérateur dont le mécanisme se réduit à acquérir l'élément suivant (si il existe). Ce mécanisme empêche l'utilisation directe de certaines fonctionnalités mentionnées dans la section « DataLoader », comme le shuffling et le mutiprocessing, qui se basent sur des manipulations d'indices et ont besoin d'une vision globale de la base de données.

Datasets​

Datasets prédéfinis dans PyTorch​

PyTorch propose un ensemble de Datasets prédéfinis dans les librairies torchvision, torchaudio et torchtext. Ces librairies gèrent la création d'un objet Dataset pour des bases de données standards listées dans les documentations officielles :

Le chargement d'une base données se fait via le module Datasets. Par exemple, le chargement de la base de données d'images ImageNet peut se faire avec torchvision de la manière suivante :

import torchvision

# load imagenet dataset stored in DSDIR
root = os.environ['DSDIR']+'/imagenet'
imagenet_dataset = torchvision.datasets.ImageNet(root=root)

La plupart du temps, il est possible de différencier au chargement les données dédiées à l'entraînement des données dédiées à la validation. Par exemple, pour la base ImageNet :

import torchvision

# load imagenet dataset stored in DSDIR
root = os.environ['DSDIR']+'/imagenet'

## load data for training
imagenet_train_dataset = torchvision.datasets.ImageNet(root=root,
                                                       split='train')
## load data for validation
imagenet_val_dataset = torchvision.datasets.ImageNet(root=root,
                                                     split='val')

Chaque fonction de chargement propose ensuite des fonctionnalités spécifiques aux bases de données (qualité des données, extraction d'une sous-partie des données, etc). Nous vous invitons à consulter les documentations officielles pour plus de détails.

remarque

La librairie torchvision contient une fonction générique de chargement torchvision.Datasets.ImageFolder. Elle est adaptée à toute base de données d’images, sous condition que celle-ci soit stockée dans un certain format (voir la documentation officielle pour plus de détails).

important

Certaines fonctions proposent de télécharger les bases données en ligne grâce à l’argument download=True. Nous vous rappelons que les nœuds de calcul Jean Zay n’ont pas accès à internet et que de telles opérations doivent se faire en amont depuis une frontale ou un nœud de pré/post-traitement. Nous vous rappelons également que des bases de données publiques et volumineuses sont déjà disponibles sur l’espace commun DSDIR de Jean Zay. Cet espace peut-être enrichi sur demande auprès de l’assistance IDRIS (assist@idris.fr).

Datasets personnalisés​

Il est possible de créer ses propres classes Datasets en définissant trois fonctions caractéristiques :

  • __init__ initialise la variable contenant les données à traiter
  • __len__ retourne la longueur de la base de données
  • __getitem__ retourne la donnée correspondant à un indice donnée

Par exemple :

class myDataset(Dataset):

    def __init__(self, data):
    # Initialise dataset from source dataset
    self.data = data

    def __len__(self):
    # Return length of the dataset
    return len(self.data)

    def __getitem__(self, idx):
    # Return one element of the dataset according to its index
    return self.data[idx]

Transformations​

Transformations prédéfinies dans PyTorch​

Les librairies torchvision, torchtext et torchaudio offrent un panel de transformations pré-implémentées, accessibles via le module transforms de la classe Datasets. Ces transformations sont listées dans les documentations officielles :

Les instructions de transformation sont portées par l'objet Dataset. Il est possible de cumuler différents types de transformations grâce à la fonction transforms.Compose(). Par exemple, pour redimensionner l'ensemble des images de la base de données ImageNet :

import torchvision

# define list of transformations to apply
data_transform = torchvision.transforms.Compose([torchvision.transforms.Resize((300,300)),
                                                 torchvision.transforms.ToTensor()])

# load imagenet dataset and apply transformations
root = os.environ['DSDIR']+'/imagenet'
imagenet_dataset = torchvision.datasets.ImageNet(root=root,
                                                 transform=data_transform)
remarque

La transformation transforms.ToTensor() permet de convertir une image PIL ou un tableau NumPy en tenseur.

Pour appliquer des transformations sur un Dataset personnalisé, il faut modifier celui-ci en conséquence, par exemple de la manière suivante :

class myDataset(Dataset):

    def __init__(self, data, transform=None):
        # Initialise dataset from source dataset
        self.data = data
        self.transform = transform

    def __len__(self):
        # Return length of the dataset
        return len(self.data)

    def __getitem__(self, idx):
        # Return one element of the dataset according to its index
        x = self.data[idx]

        # apply transformation if requested
        if self.transform:
            x = self.transform(x)

        return x

Transformations personnalisées​

Il est aussi possible de créer ses propres transformations en définissant des fonctions callable et en les communiquant directement à transforms.Compose(). On peut par exemple définir des transformations de type somme (Add) et multiplication (Mult) de la manière suivante :

# define Add tranformation
class Add(object):
    def __init__(self, value):
        self.value = value

    def __call__(self, sample):
        # add a constant to the data
        return sample + self.value

# define Mult transformation
class Mult(object):
    def __init__(self, value):
        self.value = value

    def __call__(self, sample):
        # multiply the data by a constant
        return sample * self.value

# define list of transformations to apply
data_transform = transforms.Compose([Add(2),Mult(3)])

DataLoaders​

Un objet DataLoader est une sur-couche d'un objet Dataset qui permet de structurer les données (création de batches), de les pré-traiter (shuffling, transformations) et de les diffuser aux GPU pour la phase d'entraînement.

Le DataLoader est un objet de la classe torch.utils.data.DataLoader :

import torch

# define DataLoader for a given dataset
dataloader = torch.utils.data.DataLoader(dataset)

Optimisation des paramètres de chargement des données​

Les arguments paramétrables de la classe DataLoader sont les suivants :

DataLoader(dataset,
           shuffle=False,
           sampler=None, batch_sampler=None, collate_fn=None,
           batch_size=1, drop_last=False,
           num_workers=0, worker_init_fn=None, persistent_workers=False,
           pin_memory=False, timeout=0,
           prefetch_factor=2, *
          )

Traitement aléatoire des données d'entrée​

L'argument shuffle=True permet d'activer le traitement aléatoire des données d'entrée. Attention, cette fonctionnalité doit être déléguée au sampler si vous utilisez un sampler distribué (voir point suivant).

Distribution des données sur plusieurs processus en vue d'un apprentissage distribué​

L'argument sampler permet de spécifier le type d'échantillonnage de la base de données que vous souhaitez mettre en œuvre. Pour distribuer les données sur plusieurs processus, il faut utiliser le sampler DistributedSampler fourni par la classe torch.utils.data.distributed de PyTorch. Par exemple :

import idr_torch # IDRIS package available in all PyTorch modules

# define distributed sampler
data_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset,
                                                               shuffle=True,
                                                               num_replicas=idr_torch.size,
                                                               rank=idr_torch.rank)

Ce sampler prend en argument l'ordre d'activation du shuffling, le nombre de processus disponibles num_replicas et le rang local rank. L'étape de shuffling est déléguée au sampler pour pouvoir être traitée en parallèle. Le nombre de processus et le rang local sont déterminés à partir de l'environnement Slurm dans lequel le script d'entraînement a été lancé. On utilise ici la librairie idr_torch pour récupérer ces informations. Cette librairie est développée par l'IDRIS et est présente dans l'ensemble des modules PyTorch sur Jean Zay.

remarque

Le sampler DistributedSampler est adapté à la stratégie de parallélisme de données torch.nn.parallel.DistributedDataParallel que nous documentons sur cette page.

Optimisation de l'utilisation des ressources lors de l'apprentissage​

La taille de batch est définie par l'argument batch_size. Une taille de batch est optimale si elle permet une bonne utilisation des ressources de calcul, c'est-à-dire si la mémoire de chaque GPU est sollicitée au maximum et que la charge de travail est répartie équitablement entre les GPU.

Il arrive que la quantité de données d'entrée ne soit pas un multiple de la taille de batch demandée. Dans ce cas, pour éviter que le DataLoader ne génère un batch « incomplet » avec les dernières données extraites, et ainsi éviter un déséquilibre de la charge de travail entre GPU, il est possible de lui ordonner d'ignorer ce dernier batch avec l'argument drop_last=True. Cela peut néanmoins représenter une perte d'information qu'il faut avoir estimée en amont.

Recouvrement transfert/calcul​

Il est possible d'optimiser les transferts de batches du CPU vers le GPU en générant du recouvrement transfert/calcul.

Une première optimisation consiste à pré-charger les prochains batches à traiter pendant l'entraînement. La quantité de batches pré-chargés est contrôlée par l'argument prefetch_factor. Par défaut, cette valeur est fixée à 2, ce qui convient dans la plupart des cas.

Une deuxième optimisation consiste à demander au DataLoader de stocker les batches en mémoire épinglée (pin_memory=True) sur le CPU. Cette stratégie permet d'éviter certaines étapes de recopie lors des transferts du CPU vers le GPU. Elle permet également d'utiliser le mécanisme d'asynchronisme non_blocking=True lors d'appels aux fonctions de transfert .to() ou .device().

Accélération du pré-traitement des données (transformations)​

Le pré-traitement des données (transformations) est une étape gourmande en ressources CPU. Pour l'accélérer, il est possible de paralléliser les opérations sur plusieurs CPU grâce à la fonctionnalité de multiprocessing du DataLoader. Le nombre de processus impliqués est spécifié par l'argument num_workers.

L'argument persistent_workers=True permet de maintenir les processus actifs tout au long de l'entraînement, évitant ainsi leurs réinitialisations à chaque epoch. Ce gain de temps implique en contrepartie une occupation de la mémoire RAM potentiellement importante, surtout si plusieurs DataLoaders sont utilisés.

Exemple complet de chargement optimisé de données​

Voici un exemple complet de chargement de données optimisé de la base de données ImageNet pour un apprentissage distribué sur Jean Zay :

import torch
import torchvision
import idr_torch # IDRIS package available in all PyTorch modules

# define list of transformations to apply
data_transform = torchvision.transforms.Compose([torchvision.transforms.Resize((300,300)),
                                                 torchvision.transforms.ToTensor()])

# load imagenet dataset and apply transformations
root = os.environ['DSDIR']+'/imagenet'
imagenet_dataset = torchvision.datasets.ImageNet(root=root,
                                                 transform=data_transform)

# define distributed sampler
data_sampler = torch.utils.data.distributed.DistributedSampler(imagenet_dataset,
                                                               shuffle=True,
                                                               num_replicas=idr_torch.size,
                                                               rank=idr_torch.rank
                                                              )

# define DataLoader
batch_size = 128                       # adjust batch size according to GPU type (16GB or 32GB in memory)
drop_last = True                       # set to False if it represents important information loss
num_workers = 4                        # adjust number of CPU workers per process
persistent_workers = True              # set to False if CPU RAM must be released
pin_memory = True                      # optimize CPU to GPU transfers
non_blocking = True                    # activate asynchronism to speed up CPU/GPU transfers
prefetch_factor = 2                    # adjust number of batches to preload

dataloader = torch.utils.data.DataLoader(imagenet_dataset,
                                         sampler=data_sampler,
                                         batch_size=batch_size,
                                         drop_last=drop_last,
                                         num_workers=num_workers,
                                         persistent_workers=persistent_workers,
                                         pin_memory=pin_memory,
                                         prefetch_factor=prefetch_factor
                                        )

# loop over batches
for i, (images, labels) in enumerate(dataloader):
    images = images.to(gpu, non_blocking=non_blocking)
    labels = labels.to(gpu, non_blocking=non_blocking)

Mise en pratique sur Jean Zay​

Pour mettre en pratique la documentation ci-dessus et vous faire une idée des gains apportés par chacune des fonctionnalités proposées par le DataLoader PyTorch, vous pouvez récupérer le Notebook notebook_data_preprocessing_pytorch.ipynb dans le DSDIR. Par exemple, pour le récupérer dans votre WORK :

cp $DSDIR/examples_IA/Torch_parallel/notebook_data_preprocessing_pytorch.ipynb $WORK

Vous pouvez ensuite exécuter le Notebook à partir d'une machine frontale de Jean Zay (voir notre documentation sur l'accès à JupyterHub pour en savoir plus sur l'usage des Notebooks sur Jean Zay).

Documentation officielle​

Your opinion matters!

To give your feedback, report an error, or suggest an improvement, click here:

quick anonymous questionnaire

This questionnaire is temporary and will take less than a minute, so take the opportunity!