Comment accélérer les modèles PyTorch : guide pratique torch.profiler

PyTorch inclut l'outil intégré torch.profiler pour analyser les performances des modèles. Il suit chaque opération sur GPU et CPU, aidant à identifier les goulots d'étranglement dans les calculs et à optimiser le temps d'entraînement.

Pourquoi profiler les modèles

Sans profilage, le développeur ne voit que le chiffre final : « l'époque a pris 2 heures ». Mais pourquoi 2 heures ? Où a disparu le temps ? Peut-être que le dataloader est lent ? Ou l'opération sur GPU est inefficace ? Ou le GPU reste inactif en attendant les données du CPU ? Les tentatives d'optimisation à l'aveugle, c'est deviner. Vous changez le batch size, la vitesse de chargement, la précision des calculs, mais le résultat ne s'améliore pas parce que vous n'optimisez pas le bon goulot d'étranglement. torch.profiler vous épargne cette perte de temps : il montre la distribution exacte du temps et de la mémoire entre les opérations.

Comment fonctionne torch.profiler

L'outil suit l'exécution du code au niveau des noyaux CUDA et des threads CPU. Pour chaque opération — par exemple, une multiplication matricielle dans une couche Linear ou une convolution dans Conv2d — il enregistre les heures de début et de fin, ainsi que la mémoire utilisée. Les résultats peuvent être exportés dans un format compatible avec Chrome DevTools et visualisés sous forme de timeline.

Les principales métriques du profilage :

• Exclusive time — le temps que l'opération a pris seule (sans opérations imbriquées)
• Inclusive time — le temps de l'opération y compris tous les appels imbriqués
• Memory — l'utilisation maximale de la mémoire GPU pendant l'opération
• Sync time — le temps consacré à la synchronisation entre CPU et GPU (goulot d'étranglement numéro un)

Quand torch.profiler révèle les problèmes

Scénario typique : le modèle s'entraîne lentement, vous pensez qu'il faut une GPU plus rapide, mais le profilage montre que 60 % du temps de l'époque est consacré à l'attente des données du dataloader. Le problème n'est pas dans le GPU, mais dans le chargement des données. Solution : augmentez num_workers dans DataLoader ou utilisez pinned memory.

Autre exemple : le modèle s'entraîne sur un système distribué (plusieurs GPU), et le profilage montre que beaucoup de temps est consacré à la synchronisation des gradients entre les appareils. C'est un signal pour optimiser le graphe de communication ou reconsidérer la stratégie de parallélisme.

Point important : le profilage ajoute lui-même une surcharge. Si vous lancez le profiler sur tout le cycle d'entraînement, il peut ralentir les calculs de 10 à 30 %. Pour des mesures précises, il est préférable de profiler uniquement le bloc de code qui vous intéresse (par exemple, une seule passe forward-backward).

Ce que cela signifie

torch.profiler est la première étape avant toute optimisation. Ne devinez pas où se trouve le goulot d'étranglement : profilez, regardez les résultats, et optimisez exactement ce qui ralentit. Cela vous fera économiser des semaines d'expérimentation et dirigera vos efforts là où ils apporteront vraiment des résultats. Pour les ingénieurs ML qui souhaitent accélérer l'entraînement des modèles, c'est un outil fondamental.