PyTorch Profiler : lire les coûts et exporter une trace CPU

PyTorch Profiler enregistre les opérations dans une partie choisie de votre programme. Il permet de repérer les opérateurs coûteux, de les relier à des étapes nommées du code et d’examiner leur ordre d’exécution. Commencez par une question précise : quelles opérations occupent le temps de ce passage avant ?

Ce guide exécute une petite charge d’inférence sur CPU, capture deux lots et exporte une trace ainsi que des tableaux de temps et de mémoire. L’exemple a été testé sous Windows avec Python 3.12.14 et PyTorch 2.15.0+cpu. Il utilise des entrées synthétiques, sans GPU ni téléchargement de modèle ou de données. Les mesures servent à apprendre à lire une capture ; elles ne constituent pas un benchmark du modèle.

Définir le travail à mesurer

Le test utilise un réseau à deux couches linéaires séparées par une ReLU. Chaque lot contient 128 lignes et 256 caractéristiques d’entrée ; la sortie en contient 64. Le modèle et les entrées sont créés avant le profilage. La capture porte donc sur des passages avant répétés, sans mesurer le démarrage, le chargement des données ou une requête applicative complète.

Le script active le mode évaluation avec model.eval() et s’exécute dans torch.inference_mode(). Choisissez les modes correspondant au travail à mesurer. Pour profiler un entraînement, il faut capturer son déroulement réel, dont la rétropropagation et les mises à jour de l’optimiseur. Notre guide sur eval, no_grad et inference_mode (en anglais) explique ces différences.

Un seul thread CPU rend la configuration du test explicite. Ce réglage ne recommande pas de réduire le nombre de threads de votre application. Notez vos paramètres de threads, la forme des lots, le mode du modèle, le périphérique et la version de PyTorch lors de chaque capture.

Exécuter une capture CPU planifiée

Enregistrez le code suivant dans profile_cpu.py, dans un environnement où PyTorch est installé. Il crée un nouveau répertoire et refuse d’écraser une exécution existante. La fonction de rappel sauvegarde les deux tableaux et la trace de la capture terminée.

import json
from pathlib import Path

import torch
from torch import nn
from torch.profiler import ProfilerActivity, profile, record_function, schedule


def run_profile(output_dir):
    output_dir = Path(output_dir)
    output_dir.mkdir(parents=True, exist_ok=False)
    torch.manual_seed(7)
    torch.set_num_threads(1)  # A fixed lab setting, not a tuning recommendation.
    model = nn.Sequential(nn.Linear(256, 512), nn.ReLU(), nn.Linear(512, 64))
    model.eval()
    inputs = torch.randn(128, 256)
    captures = []

    def save_trace(prof):
        stem = output_dir / f"cpu-step-{prof.step_num}"
        averages = prof.key_averages()
        time_table = averages.table(sort_by="self_cpu_time_total", row_limit=12)
        memory_table = averages.table(sort_by="self_cpu_memory_usage", row_limit=12)
        stem.with_suffix(".time.txt").write_text(time_table, encoding="utf-8")
        stem.with_suffix(".memory.txt").write_text(memory_table, encoding="utf-8")
        prof.export_chrome_trace(str(stem.with_suffix(".json")))
        captures.append({
            "step": prof.step_num,
            "operators": [{"name": e.key, "calls": e.count,
                           "self_cpu_us": e.self_cpu_time_total,
                           "total_cpu_us": e.cpu_time_total,
                           "self_cpu_bytes": e.self_cpu_memory_usage}
                          for e in averages],
        })
        print(time_table)
        print(memory_table)

    with torch.inference_mode():
        reference = model(inputs)
        for _ in range(3):
            model(inputs)  # Warm up the workload before starting the profiler.
        with profile(
            activities=[ProfilerActivity.CPU],
            schedule=schedule(wait=1, warmup=1, active=2, repeat=1),
            on_trace_ready=save_trace,
            profile_memory=True,
            record_shapes=False,
            with_stack=False,
        ) as prof:
            for step in range(6):
                with record_function(f"batch_{step}"):
                    with record_function("forward"):
                        output = model(inputs)
                    del output
                prof.step()  # Advance once after each logical batch.
        torch.testing.assert_close(model(inputs), reference, rtol=0, atol=0)

    report = {"torch": torch.__version__, "device": "cpu",
              "threads": torch.get_num_threads(), "input_shape": list(inputs.shape),
              "output_shape": list(reference.shape), "captures": captures}
    (output_dir / "summary.json").write_text(json.dumps(report, indent=2), encoding="utf-8")
    return report


if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument("output_dir", help="New directory for this capture")
    args = parser.parse_args()
    run_profile(args.output_dir)

Exécutez-le avec un nom de répertoire encore inutilisé :

python profile_cpu.py profiler-run-01

Le répertoire contient cpu-step-4.json, cpu-step-4.time.txt, cpu-step-4.memory.txt et summary.json. Le script compare aussi une prédiction avant et après le profilage pour vérifier que le résultat est inchangé. Il ne sauvegarde pas les poids du modèle et ne modifie aucun jeu de données.

Vérifier les itérations enregistrées

Le script effectue deux préparations distinctes. Trois passages avant non enregistrés préparent la charge avant le démarrage du profiler. Dans le profiler, la planification attend une itération, en utilise une pour sa préparation, puis en enregistre deux. La boucle exécute tout de même les six itérations.

• Lot 0 : attente, sans plage de lot capturée.
• Lot 1 : préparation du profiler, exclue de la capture exportée.
• Lots 2 et 3 : capture active.
• Lots 4 et 5 : aucune nouvelle capture, car repeat=1.

Notre trace contenait exactement batch_2 et batch_3, avec deux plages forward. La fonction de rappel s’est exécutée une seule fois, après le quatrième appel à prof.step(), d’où le nom du fichier. Placez cet appel à la fin de chaque étape logique. Si une étape correspond à une mise à jour de l’optimiseur plutôt qu’à un lot, adaptez les libellés et l’avancement. La référence de l’API Profiler décrit la planification et les fonctions de rappel.

Lire le temps propre avant d’additionner les totaux

Le tableau regroupe les appels répétés portant le même nom d’opérateur. Self CPU exclut le temps des appels imbriqués ; CPU total les inclut. CPU time avg correspond au total divisé par le nombre d’appels. La recette officielle du profiler explique ces colonnes et le suivi de mémoire.

Ces lignes sélectionnées proviennent de notre capture, en microsecondes. Elles ont été remises en forme à partir de summary.json ; les temps varieront lors d’une autre exécution.

Name                Calls   Self CPU (us)   CPU total (us)
forward                 2           358.1           3085.1
aten::linear            4           265.3           2455.3
aten::addmm             4          1522.5           1579.8
aten::clamp_min         2            40.6             40.6

Les deux passages avant ont produit quatre appels linéaires, puisque le modèle comporte deux couches linéaires. Dans cette capture, aten::addmm présentait le temps propre le plus élevé. Les opérations matricielles constituent donc une piste à examiner pour ce cas précis. Cela ne prouve pas que la multiplication matricielle limite tous les modèles.

Les totaux de forward, aten::linear et aten::addmm se chevauchent, car ces opérations sont imbriquées. Les additionner compterait une partie du travail plusieurs fois. Un total élevé associé à un temps propre faible peut aussi désigner une opération qui encadre des appels enfants coûteux. Examinez cette relation dans la trace avant de modifier le code.

La plage nommée forward inclut également le travail Python entre ses opérations enfants. Son temps propre ne correspond pas à une couche mathématique distincte. Pour un exemple aussi petit, le coût du profiler et des libellés peut représenter une part substantielle du temps observé.

Interpréter les allocations de mémoire

Avec profile_memory=True, le second tableau est trié par self_cpu_memory_usage. Les valeurs positives et négatives décrivent les allocations et les libérations de tenseurs attribuées au travail enregistré. La mémoire propre exclut les appels enfants. Ces colonnes ne donnent pas le pic de mémoire résidente du processus.

Dans ce test, aten::addmm indiquait 589 824 octets d’allocation propre : deux passages capturés × 128 lignes × (512 + 64) sorties × quatre octets par valeur float32. L’allocation de la ReLU figurait séparément. Les tenseurs intermédiaires et les sorties finales étaient libérés pendant chaque lot ; les plages des lots présentaient donc une variation nette nulle, malgré les allocations réalisées.

Une variation nette nulle ne signifie pas qu’aucune mémoire n’était nécessaire. Les paramètres, les entrées et la prédiction de référence existaient avant la fenêtre active. Cette capture ne permet pas de reconstruire leurs allocations antérieures ni d’établir le budget mémoire complet du processus.

Ouvrir la trace et suivre un lot

Ouvrez cpu-step-4.json dans le visualiseur Perfetto avec l’option d’ouverture d’un fichier local. Le guide de l’interface Perfetto décrit le chargement et la navigation. Recherchez batch_2, examinez sa plage forward, puis suivez les opérations linéaires et la ReLU imbriquées. Comparez avec batch_3.

Le tableau classe le travail agrégé ; la chronologie aide à situer son exécution. Nos contrôles ont analysé le JSON exporté et vérifié les noms de lots, les nombres d’appels, les durées et les événements d’allocation et de libération. Ils n’ont pas inclus d’inspection visuelle dans le visualiseur.

Modifier un point et le mesurer séparément

Utilisez la capture pour formuler une explication vérifiable, puis comparez le code initial et le code modifié sans profiler actif. Conservez les mêmes entrées, modes, paramètres de threads et périphérique, préparez les deux versions et répétez les mesures. torch.utils.benchmark.Timer prend en charge la préparation et les mesures répétées. Vérifiez aussi la justesse des sorties.

Pour un second passage de diagnostic, activez record_shapes=True si la taille des entrées est en cause, puis regroupez les moyennes avec group_by_input_shape=True. L’enregistrement des formes et des piles d’appels ajoute un coût ; celui des formes conserve aussi temporairement des références aux tenseurs. Laissez ces options désactivées lorsqu’elles ne contribuent pas au diagnostic.

Si le délai se produit avant le passage avant, élargissez la zone capturée à la récupération et à la préparation des lots. Notre guide DataLoader (en anglais) traite ces parties du parcours des données. Le script utilise un tenseur déjà créé : il ne mesure ni lecture de fichiers, ni processus de prétraitement, ni transfert du CPU vers un accélérateur.

Ce que vérifie cet exemple CPU

Dix contrôles locaux ont couvert la fenêtre de capture, le nombre de rappels et d’opérateurs, les temps imbriqués, les allocations et libérations, les fichiers exportés, la conservation du résultat et le refus d’écraser une capture. Le JSON ne contenait aucun événement de kernel d’accélérateur. L’environnement sans NumPy a émis un avertissement d’initialisation, mais le script a abouti en utilisant uniquement des tenseurs.

Profiler les kernels GPU exige un périphérique, une installation et des fonctions de profilage adaptés. Ajouter une activité CUDA à un exemple CPU ne fournit pas ces preuves. Ce test valide la capture et l’export de l’inférence CPU ; il ne démontre ni limitation GPU, ni performance d’entraînement, ni comportement distribué, ni débit sur une VM cloud, ni accélération.