Jean Zay : Outils de profilage d'applications TensorFlow et PyTorch

Le profilage est une étape indispensable de l'optimisation d'un code. Son but est de cibler les étapes d'exécution les plus coûteuses en temps ou mémoire, et de visualiser la répartition de la charge de travail entre GPU et CPU.

Profiler une exécution est une opération très gourmande en temps. Cette étude se fait donc en général sur une ou deux itérations de votre entraînement.

Profilage d'applications TensorFlow

TensorFlow intègre la fonctionnalité de profilage TensorFlow Profiler.

Le TensorFlow Profiler requiert des versions de TensorFlow et TensorBoard supérieures ou égales à 2.2. Sur Jean-Zay, il est disponible sous les versions 2.2.0 et supérieures de TensorFlow, en chargeant le module adéquat. Par exemple :

$ module load tensorflow-gpu/py3/2.2.0

Instrumenter un code TensorFlow pour le profilage

Le code doit intégrer un rappel TensorBoard, comme expliqué dans la page Outils de visualisation TensorBoard pour TensorFlow et PyTorch.

Visualisation du profil d'un code TensorFlow

La visualisation des résultats de TensorFlow Profiler est possible via TensorBoard, dans l'onglet PROFILE. L'accès à TensorBoard est documenté ici.

L'onglet PROFILE de TensorBoard s'ouvre sur l'Overview Page. Elle permet d'avoir un résumé de la performance en temps d'exécution des différentes étapes de calcul. Cela permet de savoir rapidement qui de l'apprentissage, du chargement des données ou du préprocessing des données, est le plus consommateur en temps d'exécution.

Puis la page Trace Viewer permet d'avoir une vue plus détaillée de la séquence d'exécutions, en distinguant les opérations exécutées sur GPU et sur CPU. Par exemple :

Dans cet exemple, on voit que le GPU (en haut) est peu utilisé par rapport au CPU (en bas) la plupart du temps. Les blocs de couleurs montrent que le GPU est utilisé seulement à la fin des steps, alors que le CPU est utilisé régulièrement sur certain threads. Une optimisation est certainement possible en répartissant mieux le travail entre GPU et CPU.

Profilage d'applications PyTorch

PyTorch intègre la fonctionnalité de profilage PyTorch Profiler.

Instrumenter un code PyTorch pour le profilage

Dans le code PyTorch, il faut:

  • importer le profiler
import torch.autograd.profiler as profiler
  • puis invoquer le profiler pendant l'exécution de la fonction d'apprentissage
with profiler.profile(use_cuda=True) as prof:
    with profiler.record_function("training_function"):
        train()

Remarques :

  • use_cuda=True : permet de récupérer les events CUDA (liés aux GPUs)
  • with profiler.record_function(“$NAME”) permet de poser un décorateur (une balise associée à un nom) pour un bloc de fonctions. Il est donc aussi intéressant de poser des décorateurs dans la fonction d'apprentissage pour des ensembles de sous-fonctions. Par exemple:
def train():
    for epoch in range(1, num_epochs+1):
        for i_step in range(1, total_step+1):
            # Obtain the batch.
            with profiler.record_function("load input batch"):
                images, captions = next(iter(data_loader))
            ...
            with profiler.record_function("Training step"):
              ...
                loss = criterion(outputs.view(-1, vocab_size), captions.view(-1))
              ...
  • à partir de la version 1.6.0 de PyTorch, il est possible de profiler l'empreinte mémoire CPU et GPU en ajoutant le paramètre profile_memory=True sous profiler.profile.

Visualisation du profilage d'un code PyTorch

Il n'est pas encore possible dans les versions actuelles de visualiser les résultats de PyTorch Profiler avec la fonctionnalité PROFILE de TensorBoard. Nous vous proposons ci-dessous d'autres solutions documentées dans Pytorch.

Visualisation d'un tableau de profilage

Une fois la fonction d'apprentissage exécutée, pour afficher les résultats du profile, il faut lancer la ligne :

print(prof.key_averages().table(sort_by="cpu_time", row_limit=10))

On obtient la liste de toutes les fonctions balisées automatiquement ou par nos soins (via les décorateurs), triée par ordre décroissant en temps CPU total. Par exemple :

|-----------------------------------   ---------------  ---------------  -------------  -------------- 
Name                                   Self CPU total   CPU total        CPU time avg  Number of Calls 
|-----------------------------------   ---------------  ---------------  -------------  -------------- 
training_function                     1.341s            62.089s          62.089s          1                
load input batch                      57.357s           58.988s          14.747s          4                
Training step                         1.177s            1.212s           303.103ms        4                
EmbeddingBackward                     51.355us          3.706s           231.632ms        16               
embedding_backward                    30.284us          3.706s           231.628ms        16               
embedding_dense_backward              3.706s            3.706s           231.627ms        16               
move to GPU                           5.967ms           546.398ms        136.599ms        4                
stack                                 760.467ms         760.467ms        95.058ms         8                
BroadcastBackward                     4.698ms           70.370ms         8.796ms          8                
ReduceAddCoalesced                    22.915ms          37.673ms         4.709ms          8                
|-----------------------------------  ---------------   ---------------  ---------------  ------------

La colonne Self CPU total correspond au temps passé dans la fonction elle-même, pas dans ses sous-fonctions.

La colonne Number of Calls contient le nombre de GPU utilisé par une fonction.

Ici, on voit que le temps de chargement des images (étape load input batch) est beaucoup plus important que celui de l'apprentissage du réseau de neurones (étape Training step). Le travail d'optimisation doit donc cibler le chargement des batchs.

Visualisation du profilage avec l'outil de trace Chromium

Pour afficher un Trace Viewer équivalent à celui de TensorBoard, vous pouvez également générer un fichier de trace json avec la ligne suivante :

prof.export_chrome_trace("trace.json")

Ce fichier de trace est visualisable sur l'outil de trace du projet Chromium. À partir d'un navigateur CHROME (ou CHROMIUM), il faut lancer la commande suivante dans la barre URL:

about:tracing    

Ici, on voit distinctement l'utilisation des CPU et des GPU, comme avec TensorBoard. En haut, sont représentées les fonctions CPU et en bas, sont représentées les fonctions GPU. On a 5 tâches CPU et 4 tâches GPU. Chaque bloc de couleur représente une fonction ou une sous-fonction. Nous sommes à la fin d'un chargement d'un batch de données, suivi d'une itération d'apprentissage. Dans la partie inférieure, nous voyons le forward, le backward et la synchronisation exécutés en multi-GPU.

Documentations officielles