Définir des modèles avec LightningModule

Modèles d'IA évolutifs avec PyTorch Lightning

Sergiy Tkachuk

Director, GenAI Productivity

LightningModule à la loupe

Encapsule l’architecture du modèle
Regroupe la logique d’entraînement en une unité gérable
Un plan qui apporte ordre et clarté aux projets de deep learning

Schéma PyTorch LightningModule

Définir la méthode init

Tâches clés :

Initialisation du modèle
super() :
- Gestion automatique des boucles d’entraînement
- Journalisation
- Sauvegardes (checkpointing)
Couches définies après l’initialisation
Modulaire et facile à maintenir

import lightning.pytorch as pl
import torch.nn as nn

class ClassificationModel(pl.LightningModule):
    def __init__(self, input_dim,
                 hidden_dim, num_class):
          # Initialize parent class
        super().__init__()

        # First layer
        self.layer1 = nn.Linear(input_dim,
                                hidden_dim)
        # Activation function
        self.relu = nn.ReLU()
        # Output layer
        self.layer2 = nn.Linear(hidden_dim,
                                num_class)

Implémenter la méthode forward

Étapes clés :

Définir le flux de données dans le réseau
Traiter l’entrée séquentiellement
- Transformation linéaire
- Activation
- Dernière couche et sortie

import lightning.pytorch as pl
import torch.nn as nn

class ClassificationModel(pl.LightningModule):
    def __init__(self, input_dim,
                 hidden_dim, num_class):
          ...

    def forward(self, x):

        x = self.layer1(x) # Pass input

        x = nn.ReLU(x) # Apply activation

        x = self.layer2(x) # Compute output

        return x # Return result

Exemple : classification de chiffres manuscrits

import lightning.pytorch as pl
from torch.utils.data import DataLoader
from torchvision.datasets import MNIST
from torchvision import transforms

transform = transforms.ToTensor()
train_ds = MNIST(root='.', train=True,  download=True, transform=transform)
test_ds  = MNIST(root='.', train=False, download=True, transform=transform)

train_loader = DataLoader(train_ds, batch_size=64, shuffle=True)
test_loader  = DataLoader(test_ds,  batch_size=64)


model = ClassificationModel(input_dim=28*28, hidden_dim=128, num_class=10)


trainer = pl.Trainer(max_epochs=3, accelerator='auto')
trainer.fit(model, train_loader, test_loader)

Intégrer le modèle aux tâches de classification

Se concentrer sur un cas de classification
Tout le flux dans LightningModule
Sortir les logits pour la softmax
Intégration avec Lightning Trainer

class ClassificationModel(pl.LightningModule):
  def __init__(self, input_dim, 
               hidden_dim, output_dim):
    super().__init__()

    self.hid = nn.Linear(input_dim, hidden_dim)
    self.out = nn.Linear(hidden_dim, output_dim)


  def forward(self, x):

    x = self.hidden(x)
    x = nn.ReLU(x)
    x = self.output(x)

    return x

Passons à la pratique !

Modèles d'IA évolutifs avec PyTorch Lightning