PyTorch Lightning et best practices : Code professionnel et production-ready

11 min de lecture 2217 mots

tl;dr: Guide complet PyTorch Lightning : LightningModule pour structure claire, Trainer avec multi-GPU/TPU automatique, callbacks (EarlyStopping, Checkpointing), logging professionnel (wandb, TensorBoard), et architecture de projet production-ready. Simplifie le code tout en gardant la flexibilité PyTorch.

PyTorch Lightning est un framework high-level qui structure votre code PyTorch pour le rendre propre, maintenable, et production-ready. Il automatise les détails techniques (multi-GPU, logging, checkpointing) tout en conservant la flexibilité complète de PyTorch.

Dans cet article, vous allez :

⚡ Structurer votre code avec LightningModule
🎯 Automatiser l’entraînement avec Trainer
📊 Intégrer logging (TensorBoard, wandb, MLflow)
🔔 Utiliser callbacks (EarlyStopping, Checkpointing, LR monitoring)
🏗️ Adopter une structure de projet professionnelle
🧪 Implémenter testing et reproducibilité

Niveau : Intermédiaire à Avancé

💡 Prérequis : Maîtrise de PyTorch (articles précédents). PyTorch Lightning est une surcouche : vous gardez 100% de la flexibilité PyTorch tout en gagnant en structure et automatisation.

Code et visualisation PyTorch pour PyTorch Lightning pour simplifier l’entraînement dans le développement de modèles d’IA

Pourquoi PyTorch Lightning ?

Problèmes du PyTorch Vanilla

# Code PyTorch typique : tout mélangé
for epoch in range(epochs):
    model.train()
    for batch in train_loader:
        data, target = batch
        data, target = data.cuda(), target.cuda()
        optimizer.zero_grad()
        output = model(data)
        loss = criterion(output, target)
        loss.backward()
        optimizer.step()

    # Validation
    model.eval()
    val_loss = 0
    with torch.no_grad():
        for batch in val_loader:
            # ... répéter la logique

    # Logging
    print(f"Epoch {epoch}: Loss {val_loss}")

    # Checkpointing
    if val_loss < best_loss:
        torch.save(model.state_dict(), 'best.pth')

    # Early stopping
    if patience_counter > patience:
        break

# Multi-GPU ?
if torch.cuda.device_count() > 1:
    model = nn.DataParallel(model)
    # ... modifier tout le code

# 16-bit precision ?
# ... encore plus de modifications

Problèmes :

❌ Logique métier (model, loss) mélangée avec engineering (GPU, logging)
❌ Code répétitif (train loop, val loop, test loop)
❌ Difficile à tester unitairement
❌ Multi-GPU/TPU nécessite refactoring complet
❌ Pas de standardisation entre projets

Solution : PyTorch Lightning

Principe : Séparer QUOI (logique métier) de COMMENT (infrastructure).

# Lightning : structure claire
class LitModel(L.LightningModule):
    def __init__(self):
        super().__init__()
        self.model = ...

    def forward(self, x):
        return self.model(x)

    def training_step(self, batch, batch_idx):
        x, y = batch
        y_hat = self(x)
        loss = F.cross_entropy(y_hat, y)
        return loss

    def configure_optimizers(self):
        return torch.optim.Adam(self.parameters())

# Trainer gère TOUT le reste
trainer = L.Trainer(max_epochs=10, accelerator='gpu', devices=4)
trainer.fit(model, train_loader, val_loader)

Avantages :

✅ Code 3-5× plus court
✅ Multi-GPU/TPU : changer 1 paramètre
✅ Reproductible et testable
✅ Best practices par défaut
✅ Flexibilité PyTorch complète conservée

Installation

pip install pytorch-lightning

# Avec extras pour logging
pip install pytorch-lightning[extra]  # wandb, tensorboard, etc.

LightningModule : Structure de base

Anatomie d’un LightningModule

import pytorch_lightning as L
import torch
import torch.nn as nn
import torch.nn.functional as F

class LitClassifier(L.LightningModule):
    def __init__(self, input_size, hidden_size, num_classes, learning_rate=1e-3):
        super().__init__()

        # Sauvegarder les hyperparamètres (pour checkpointing)
        self.save_hyperparameters()

        # Model
        self.layer_1 = nn.Linear(input_size, hidden_size)
        self.layer_2 = nn.Linear(hidden_size, num_classes)

    def forward(self, x):
        """Forward pass (inférence)"""
        x = F.relu(self.layer_1(x))
        x = self.layer_2(x)
        return x

    def training_step(self, batch, batch_idx):
        """
        Définir une étape d'entraînement
        Lightning gère automatiquement : backward(), optimizer.step(), etc.
        """
        x, y = batch
        y_hat = self(x)
        loss = F.cross_entropy(y_hat, y)

        # Logging
        self.log('train_loss', loss, on_step=True, on_epoch=True, prog_bar=True)

        return loss

    def validation_step(self, batch, batch_idx):
        """Étape de validation"""
        x, y = batch
        y_hat = self(x)
        loss = F.cross_entropy(y_hat, y)

        # Métriques
        acc = (y_hat.argmax(dim=1) == y).float().mean()

        self.log('val_loss', loss, prog_bar=True)
        self.log('val_acc', acc, prog_bar=True)

        return loss

    def test_step(self, batch, batch_idx):
        """Étape de test"""
        x, y = batch
        y_hat = self(x)
        loss = F.cross_entropy(y_hat, y)

        acc = (y_hat.argmax(dim=1) == y).float().mean()

        self.log('test_loss', loss)
        self.log('test_acc', acc)

    def configure_optimizers(self):
        """Définir optimizer et scheduler"""
        optimizer = torch.optim.Adam(self.parameters(), lr=self.hparams.learning_rate)

        scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
            optimizer,
            mode='min',
            factor=0.5,
            patience=5,
            verbose=True
        )

        return {
            'optimizer': optimizer,
            'lr_scheduler': {
                'scheduler': scheduler,
                'monitor': 'val_loss',  # Métrique à surveiller
            }
        }

# Instanciation
model = LitClassifier(input_size=784, hidden_size=256, num_classes=10)

Méthodes importantes

Méthode	Description	Requis
`__init__`	Initialiser modèle et hyperparams	✅
`forward`	Forward pass	✅
`training_step`	Logique d’entraînement (un batch)	✅
`validation_step`	Logique de validation	Recommandé
`test_step`	Logique de test	Optionnel
`configure_optimizers`	Définir optimizer(s) et scheduler(s)	✅
`on_train_epoch_end`	Hook fin d’epoch	Optionnel
`on_validation_epoch_end`	Hook fin de validation	Optionnel

Trainer : Automatisation complète

Le Trainer gère tout le cycle de vie de l’entraînement.

Configuration de Base

from pytorch_lightning import Trainer

trainer = Trainer(
    max_epochs=50,
    accelerator='gpu',       # 'cpu', 'gpu', 'tpu', 'mps'
    devices=2,               # Nombre de GPUs/TPUs
    precision='16-mixed',    # Mixed precision automatique
    log_every_n_steps=10,
)

# Entraîner
trainer.fit(model, train_dataloaders=train_loader, val_dataloaders=val_loader)

# Tester
trainer.test(model, dataloaders=test_loader)

# Prédire
predictions = trainer.predict(model, dataloaders=predict_loader)

Options du trainer

Accélération :

trainer = Trainer(
    accelerator='gpu',
    devices=4,                      # 4 GPUs
    strategy='ddp',                 # DistributedDataParallel
    precision='16-mixed',           # Mixed precision
    sync_batchnorm=True,            # Sync BatchNorm entre GPUs
)

Early Stopping et Checkpointing :

from pytorch_lightning.callbacks import EarlyStopping, ModelCheckpoint

checkpoint_callback = ModelCheckpoint(
    dirpath='checkpoints/',
    filename='best-{epoch:02d}-{val_loss:.2f}',
    monitor='val_loss',
    mode='min',
    save_top_k=3,               # Garder les 3 meilleurs
    save_last=True,             # Garder le dernier aussi
)

early_stop_callback = EarlyStopping(
    monitor='val_loss',
    patience=10,
    mode='min',
    verbose=True
)

trainer = Trainer(
    callbacks=[checkpoint_callback, early_stop_callback],
    max_epochs=100
)

Validation Pendant Training :

trainer = Trainer(
    val_check_interval=0.25,        # Valider 4× par epoch
    # OU
    check_val_every_n_epoch=2,      # Valider tous les 2 epochs
    limit_val_batches=0.5,          # Valider sur 50% des données val
)

Gradient Clipping et Accumulation :

trainer = Trainer(
    gradient_clip_val=0.5,          # Gradient clipping
    accumulate_grad_batches=4,      # Gradient accumulation
)

Callbacks : Modularité et réutilisabilité

Callbacks Built-in

from pytorch_lightning.callbacks import (
    EarlyStopping,
    ModelCheckpoint,
    LearningRateMonitor,
    RichProgressBar,
    TQDMProgressBar,
    ModelSummary
)

callbacks = [
    # Early stopping
    EarlyStopping(monitor='val_loss', patience=10, mode='min'),

    # Checkpointing
    ModelCheckpoint(
        monitor='val_acc',
        mode='max',
        save_top_k=1,
        filename='best-model-{epoch:02d}-{val_acc:.2f}'
    ),

    # LR monitoring
    LearningRateMonitor(logging_interval='step'),

    # Progress bar
    RichProgressBar(),

    # Model summary
    ModelSummary(max_depth=2),
]

trainer = Trainer(callbacks=callbacks)

Custom callback

from pytorch_lightning.callbacks import Callback

class PrintEpochResults(Callback):
    def on_train_epoch_end(self, trainer, pl_module):
        """Appelé à la fin de chaque epoch d'entraînement"""
        metrics = trainer.callback_metrics
        print(f"\nEpoch {trainer.current_epoch}:")
        print(f"  Train Loss: {metrics.get('train_loss', 0):.4f}")
        print(f"  Val Loss: {metrics.get('val_loss', 0):.4f}")
        print(f"  Val Acc: {metrics.get('val_acc', 0):.4f}")

class SavePredictionSamples(Callback):
    def on_validation_batch_end(self, trainer, pl_module, outputs, batch, batch_idx):
        """Sauvegarder quelques prédictions pour inspection"""
        if batch_idx == 0:  # Premier batch seulement
            x, y = batch
            y_hat = pl_module(x)
            # Sauvegarder ou logger...

Logging : TensorBoard, Weights & Biases, MLflow

TensorBoard (Built-in)

from pytorch_lightning.loggers import TensorBoardLogger

logger = TensorBoardLogger('logs/', name='my_model')

trainer = Trainer(logger=logger)
trainer.fit(model, train_loader, val_loader)

# Visualiser
# tensorboard --logdir=logs/

Weights & Biases (recommandé)

from pytorch_lightning.loggers import WandbLogger

wandb_logger = WandbLogger(
    project='pytorch-lightning-demo',
    name='experiment-1',
    log_model='all'  # Logger les checkpoints sur wandb
)

trainer = Trainer(logger=wandb_logger, max_epochs=50)
trainer.fit(model, train_loader, val_loader)

Logger des images/metrics custom :

class LitModel(L.LightningModule):
    def validation_step(self, batch, batch_idx):
        x, y = batch
        y_hat = self(x)

        # Metrics classiques
        loss = F.cross_entropy(y_hat, y)
        self.log('val_loss', loss)

        # Logger des images sur wandb
        if batch_idx == 0:
            self.logger.experiment.log({
                "examples": [wandb.Image(x[i]) for i in range(8)]
            })

        # Logger une matrice de confusion
        preds = y_hat.argmax(dim=1)
        self.logger.experiment.log({
            "confusion_matrix": wandb.plot.confusion_matrix(
                probs=None,
                y_true=y.cpu().numpy(),
                preds=preds.cpu().numpy(),
                class_names=[str(i) for i in range(10)]
            )
        })

MLflow

from pytorch_lightning.loggers import MLFlowLogger

mlf_logger = MLFlowLogger(experiment_name='lightning_exp', tracking_uri='file:./ml-runs')

trainer = Trainer(logger=mlf_logger)

Multi-Logger

# Combiner plusieurs loggers
from pytorch_lightning.loggers import TensorBoardLogger, WandbLogger

tb_logger = TensorBoardLogger('logs/')
wandb_logger = WandbLogger(project='my_project')

trainer = Trainer(logger=[tb_logger, wandb_logger])

Structure de projet professionnelle

Architecture recommandée

project/
├── README.md
├── requirements.txt
├── setup.py                 # Pour installation pip
├── .gitignore
│
├── configs/
│   ├── default.yaml         # Config par défaut
│   └── experiment_1.yaml    # Configs d'expériences
│
├── data/
│   ├── raw/                 # Données brutes
│   ├── processed/           # Données prétraitées
│   └── dataloaders.py       # DataModule
│
├── models/
│   ├── __init__.py
│   ├── classifier.py        # LightningModule
│   └── components.py        # Sous-modules réutilisables
│
├── callbacks/
│   ├── __init__.py
│   └── custom_callbacks.py
│
├── utils/
│   ├── __init__.py
│   ├── metrics.py
│   └── visualization.py
│
├── tests/
│   ├── test_model.py
│   └── test_datamodule.py
│
├── scripts/
│   ├── train.py             # Script d'entraînement
│   ├── evaluate.py          # Évaluation
│   └── predict.py           # Inférence
│
├── notebooks/
│   └── exploration.ipynb    # Analyses exploratoires
│
└── outputs/
    ├── checkpoints/
    ├── logs/
    └── predictions/

DataModule : Encapsuler les données

from pytorch_lightning import LightningDataModule
from torch.utils.data import DataLoader, random_split
from torchvision import datasets, transforms

class MNISTDataModule(LightningDataModule):
    def __init__(self, data_dir='./data', batch_size=32, num_workers=4):
        super().__init__()
        self.data_dir = data_dir
        self.batch_size = batch_size
        self.num_workers = num_workers

        self.transform = transforms.Compose([
            transforms.ToTensor(),
            transforms.Normalize((0.1307,), (0.3081,))
        ])

    def prepare_data(self):
        """Télécharger les données (appelé une seule fois)"""
        datasets.MNIST(self.data_dir, train=True, download=True)
        datasets.MNIST(self.data_dir, train=False, download=True)

    def setup(self, stage=None):
        """Créer train/val/test splits (appelé sur chaque GPU)"""
        if stage == 'fit' or stage is None:
            mnist_full = datasets.MNIST(self.data_dir, train=True, transform=self.transform)
            self.mnist_train, self.mnist_val = random_split(mnist_full, [55000, 5000])

        if stage == 'test' or stage is None:
            self.mnist_test = datasets.MNIST(self.data_dir, train=False, transform=self.transform)

    def train_dataloader(self):
        return DataLoader(
            self.mnist_train,
            batch_size=self.batch_size,
            shuffle=True,
            num_workers=self.num_workers,
            pin_memory=True
        )

    def val_dataloader(self):
        return DataLoader(
            self.mnist_val,
            batch_size=self.batch_size,
            num_workers=self.num_workers,
            pin_memory=True
        )

    def test_dataloader(self):
        return DataLoader(
            self.mnist_test,
            batch_size=self.batch_size,
            num_workers=self.num_workers,
            pin_memory=True
        )

# Utilisation
dm = MNISTDataModule(batch_size=64)
trainer.fit(model, datamodule=dm)
trainer.test(model, datamodule=dm)

Script d’entraînement (scripts/train.py)

import argparse
import pytorch_lightning as L
from pytorch_lightning.loggers import WandbLogger
from pytorch_lightning.callbacks import ModelCheckpoint, EarlyStopping

from models.classifier import LitClassifier
from data.dataloaders import MNISTDataModule

def main(args):
    # Reproducibilité
    L.seed_everything(args.seed, workers=True)

    # DataModule
    dm = MNISTDataModule(
        data_dir=args.data_dir,
        batch_size=args.batch_size,
        num_workers=args.num_workers
    )

    # Model
    model = LitClassifier(
        input_size=784,
        hidden_size=args.hidden_size,
        num_classes=10,
        learning_rate=args.lr
    )

    # Logger
    logger = WandbLogger(project='mnist-lightning', name=args.exp_name)

    # Callbacks
    checkpoint_callback = ModelCheckpoint(
        dirpath='outputs/checkpoints',
        filename='{epoch:02d}-{val_loss:.2f}',
        monitor='val_loss',
        mode='min',
        save_top_k=3
    )

    early_stop = EarlyStopping(monitor='val_loss', patience=10, mode='min')

    # Trainer
    trainer = L.Trainer(
        max_epochs=args.max_epochs,
        accelerator=args.accelerator,
        devices=args.devices,
        precision=args.precision,
        logger=logger,
        callbacks=[checkpoint_callback, early_stop],
        deterministic=True,  # Reproductibilité
    )

    # Train
    trainer.fit(model, datamodule=dm)

    # Test
    trainer.test(model, datamodule=dm, ckpt_path='best')

if __name__ == '__main__':
    parser = argparse.ArgumentParser()

    # Data
    parser.add_argument('--data_dir', type=str, default='./data')
    parser.add_argument('--batch_size', type=int, default=64)
    parser.add_argument('--num_workers', type=int, default=4)

    # Model
    parser.add_argument('--hidden_size', type=int, default=256)
    parser.add_argument('--lr', type=float, default=1e-3)

    # Training
    parser.add_argument('--max_epochs', type=int, default=50)
    parser.add_argument('--accelerator', type=str, default='gpu')
    parser.add_argument('--devices', type=int, default=1)
    parser.add_argument('--precision', type=str, default='32')
    parser.add_argument('--seed', type=int, default=42)
    parser.add_argument('--exp_name', type=str, default='experiment-1')

    args = parser.parse_args()
    main(args)

Lancement :

python scripts/train.py --batch_size=128 --hidden_size=512 --max_epochs=100 --devices=2

Testing et reproductibilité

Unit Tests

tests/test_model.py :

import torch
import pytest
from models.classifier import LitClassifier

def test_model_forward():
    """Tester le forward pass"""
    model = LitClassifier(input_size=784, hidden_size=128, num_classes=10)
    batch = torch.randn(32, 784)

    output = model(batch)

    assert output.shape == (32, 10), "Output shape incorrect"

def test_training_step():
    """Tester training_step"""
    model = LitClassifier(input_size=784, hidden_size=128, num_classes=10)
    batch = (torch.randn(32, 784), torch.randint(0, 10, (32,)))

    loss = model.training_step(batch, 0)

    assert loss is not None, "Loss should not be None"
    assert loss.ndim == 0, "Loss should be scalar"

# Lancer avec : pytest tests/

Reproductibilité

import pytorch_lightning as L

# Seed global
L.seed_everything(42, workers=True)

# Trainer avec déterminisme
trainer = L.Trainer(
    deterministic=True,  # Force deterministic algorithms
    benchmark=False,     # Disable cudnn.benchmark
)

Migration depuis PyTorch Vanilla

Avant (PyTorch)

# train.py (PyTorch vanilla)
model = Model()
optimizer = torch.optim.Adam(model.parameters())

for epoch in range(epochs):
    for batch in train_loader:
        x, y = batch
        optimizer.zero_grad()
        loss = F.cross_entropy(model(x), y)
        loss.backward()
        optimizer.step()

Après (Lightning)

# models/model.py
class LitModel(L.LightningModule):
    def __init__(self):
        super().__init__()
        self.model = Model()

    def forward(self, x):
        return self.model(x)

    def training_step(self, batch, batch_idx):
        x, y = batch
        loss = F.cross_entropy(self(x), y)
        return loss

    def configure_optimizers(self):
        return torch.optim.Adam(self.parameters())

# scripts/train.py
model = LitModel()
trainer = L.Trainer(max_epochs=10)
trainer.fit(model, train_loader)

Gains :

✅ Validation loop automatique
✅ Multi-GPU avec devices=4
✅ Checkpointing automatique
✅ Logging structuré
✅ Testing avec trainer.test()

Best Practices : Checklist

✅ Structure et organisation

Utiliser LightningModule pour tout modèle
Séparer données avec LightningDataModule
Structure de projet claire (voir ci-dessus)
Configs dans des fichiers YAML (non hardcodé)

✅ Reproductibilité

L.seed_everything() au début
Trainer(deterministic=True) si critique
Sauvegarder hyperparams avec self.save_hyperparameters()
Logger la version du code (git commit hash)

✅ Monitoring et Logging

Logger toutes les métriques importantes
Utiliser wandb ou MLflow pour expériences
ModelCheckpoint pour sauvegarder les meilleurs modèles
LearningRateMonitor pour suivre le LR

✅ Performance

precision='16-mixed' pour mixed precision
devices=N pour multi-GPU
num_workers > 0 dans DataLoader
pin_memory=True pour GPU

✅ Robustesse

EarlyStopping pour éviter overfitting
Gradient clipping si instabilité
Unit tests pour le modèle
Validation sur données séparées

Exercices Pratiques

Migrer un projet existant

Prenez un projet PyTorch existant et migrez-le vers Lightning.

Objectifs :

Convertir le modèle en LightningModule
Créer un LightningDataModule
Configurer le Trainer avec callbacks appropriés
Intégrer wandb logging
Comparer le code avant/après (nombre de lignes, clarté)

Projet complet avec best practices

Créez un projet de classification d’images from scratch avec toutes les best practices.

Objectifs :

Structure de projet professionnelle (folders, configs)
LightningModule avec augmentation de données
Callbacks : EarlyStopping, ModelCheckpoint, LearningRateMonitor
Logging avec wandb : métriques + images + confusion matrix
Unit tests pour model et datamodule
Scripts train.py, evaluate.py, predict.py
README avec instructions d’utilisation

Hyperparameter Tuning avec Lightning

Implémentez une recherche d’hyperparamètres avec Optuna ou Ray Tune.

Objectifs :

Intégrer Optuna avec PyTorchLightningPruningCallback
Optimiser : learning_rate, hidden_size, dropout, batch_size
Paralléliser les trials
Logger tous les trials sur wandb
Analyser les meilleurs hyperparams et leur impact

Conclusion

Vous maîtrisez maintenant PyTorch Lightning et les best practices :

✅ LightningModule : Structure claire et maintenable

✅ Trainer : Automatisation complète (multi-GPU, logging, checkpointing)

✅ Callbacks : Modularité et réutilisabilité

✅ Logging : TensorBoard, wandb, MLflow pour expériences

✅ DataModule : Encapsulation des données

✅ Structure de projet : Production-ready et collaborative

✅ Reproductibilité : Seed, déterminisme, versioning

Félicitations !

Vous avez complété la série PyTorch complète :

✅ Introduction et premiers tensors
✅ Tensors et opérations avancées
✅ nn.Module et architectures
✅ Training loops et optimizers
✅ Datasets et DataLoaders
✅ Transfer learning
✅ CNNs et Computer Vision
✅ RNNs, LSTMs et Transformers
✅ GANs et modèles génératifs
✅ Déploiement en production
✅ Optimisation et performance
✅ PyTorch Lightning et best practices

Vous êtes maintenant capable de :

🔥 Implémenter n’importe quelle architecture PyTorch
🚀 Entraîner sur multi-GPU avec optimisations
🏭 Déployer en production avec APIs
📊 Structurer du code professionnel et maintenable
🧪 Tester, logger, et reproduire vos expériences

Pour aller plus loin

Documentation :

PyTorch Lightning Docs - Documentation officielle complète
Lightning Examples - Exemples officiels

Communauté :

Lightning Discord - Support communautaire
Lightning Blog - Tutoriels et best practices

Frameworks Complémentaires :

Hydra : Gestion de configs avancée
Optuna : Hyperparameter optimization
DVC : Version control pour datasets et modèles

Félicitations pour avoir maîtrisé PyTorch de A à Z !

Vous êtes prêt à créer des modèles d’IA state-of-the-art et à les déployer en production. Bonne chance dans vos projets deep learning !