Table des matières
PyTorch : Chargement de bases de données pour l’apprentissage distribué d’un modèle
Dans cette page, nous mettons en pratique la gestion des Datasets et DataLoaders pour l'apprentissage distribué de modèle PyTorch. Nous nous intéressons aux problématiques présentées dans la page mère sur le chargement des données.
Nous présentons ici l'usage :
- des Datasets (prédéfinis et personnalisés)
- des outils de transformations des données d'entrée (prédéfinis et personnalisés)
- des DataLoaders
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/RawImages' 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/RawImages' ## 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 documentation officielles pour plus de détails.
Remarques :
- 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étail). - 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 tranforms 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/RawImages' 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 vu 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é 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’appel 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/RawImages' 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 Jupyter 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 également télécharger le Notebook ici.
Vous pouvez ensuite ouvrir et exécuter le Notebook depuis notre service jupyterhub. Pour l'utilisation de jupyterhub, vous pouvez consulter les documentations dédiées : Jean Zay : Accès à JupyterHub et documentation JupyterHub Jean-Zay.