PyTorch Lightning : le PyTorch boosté

  MINI-COURS #30      ⏱️ 25-30 minutes de lecture

PyTorch est la référence incontournable pour les chercheurs et praticiens qui construisent des modèles de deep learning. Sa flexibilité, son API pythonique intuitive et sa facilité de prise en main en font le choix naturel dans l'écosystème.

Créer un modèle de deep learning en PyTorch ne demande que trois étapes :

1. On définit une classe de modèle qui hérite de nn.Module :

2. On déclare tous les composants du réseau (couches, dropout, batch norm, etc.) dans la méthode __init__() :

3. On définit le passage avant (forward pass) du réseau de neurones dans la méthode forward() :

C'est tout !

Une fois le réseau défini, on peut enchaîner avec l'entraînement en déclarant l'optimiseur, la fonction de perte, etc. — sans avoir à coder explicitement le passage arrière (backward pass).

Concrètement, on peut définir la boucle d'entraînement comme ci-dessous et entraîner le modèle facilement :

Définir le réseau, c'est simple et élégant, on l'a vu.

Mais à mesure que les modèles gagnent en complexité et en taille, plusieurs défis émergent avec PyTorch :

Avec des modèles complexes, gérer manuellement la boucle d'entraînement devient fastidieux.

Itérer sur le dataset, exécuter les passes avant et arrière, mettre à jour les paramètres du modèle — tout ça reste standardisé (comme ci-dessous), mais nécessite quand même une attention constante et un effort de maintenance.

Le logging est essentiel pour suivre le processus d'entraînement et analyser les performances du modèle.

PyTorch ne fournit pas de support natif pour le logging : il faut implémenter sa propre solution ou intégrer un framework externe.

Plus les modèles grossissent, plus il devient nécessaire de les entraîner sur plusieurs GPU ou plusieurs machines pour réduire le temps d'entraînement.

PyTorch supporte l'entraînement distribué, mais l'implémentation est complexe : mise en place des processus, synchronisation des gradients, gestion de la communication entre processus.

Déboguer un entraînement distribué est particulièrement ardu, en raison de la complexité du setup et des problèmes potentiels liés à la communication inter-processus.

L'entraînement en précision mixte consiste à utiliser une précision réduite (par exemple, des flottants demi-précision) pour certaines parties du processus d'entraînement — ce qui permet de réduire l'utilisation mémoire et d'accélérer l'entraînement.

PyTorch supporte la précision mixte, mais gérer manuellement la précision de chaque opération reste un vrai défi.

On a d'ailleurs couvert ce sujet en profondeur dans un numéro récent de la newsletter :

PyTorch supporte nativement les GPU, mais faire tourner un modèle sur des TPU (Tensor Processing Units) requiert une configuration supplémentaire non triviale.

Ce constat est clair : PyTorch ne propose pas de solutions clé en main pour de nombreuses tâches importantes, ce qui génère du code répétitif (boilerplate) et multiplie les risques d'erreurs.

Ces défis ne sont pas critiques pour tous les projets, notamment les plus petits ou les plus simples.

À petite échelle, le surcoût de gestion des boucles, du logging et de l'entraînement distribué est souvent négligeable. Mais à mesure que les modèles gagnent en complexité et en taille, ces problèmes deviennent de plus en plus prononcés.

PyTorch Lightning répond à chacun des défis listés ci-dessus.

Pense à PyTorch Lightning comme à une surcouche légère autour de PyTorch qui abstrait le code répétitif qu'on écrit habituellement, et rend le processus d'entraînement plus lisible et plus fluide.

De la même façon que Keras est une surcouche sur TensorFlow, PyTorch Lightning est une surcouche sur PyTorch — mais une surcouche qui rend l'entraînement bien plus efficace que la méthode traditionnelle.

Ainsi, n'importe quel modèle PyTorch peut être utilisé comme modèle PyTorch Lightning.

Comme la bibliothèque est une surcouche optimisée autour de PyTorch, ses développeurs annoncent une réduction du code répétitif de 70 à 80 %, ce qui minimise les surfaces d'erreur et te permet de te concentrer sur ce qui compte vraiment — la valeur métier — plutôt que sur l'ingénierie.

De plus, comme on le verra, PyTorch Lightning permet de définir le modèle et la logique d'entraînement de façon claire et concise, ce qui libère du temps pour la recherche plutôt que pour les détails d'implémentation.

Son utilité parle d'elle-même : le dépôt GitHub compte plus de 31 000 étoiles :

Voici comment PyTorch Lightning adresse concrètement chaque défi évoqué plus haut :

L'un des autres grands atouts de PyTorch Lightning, c'est la minimalité de son API. Dans la grande majorité des cas, il suffit de maîtriser deux classes : LightningModule et Trainer. Le reste n'est que du PyTorch bien organisé.

Si tout ça n'est pas encore parfaitement clair, pas d'inquiétude — on va tout parcourir ensemble, étape par étape.

💡 En continuant, tu vas découvrir :