Definire modelli con LightningModule

Modelli di AI scalabili con PyTorch Lightning

Sergiy Tkachuk

Director, GenAI Productivity

LightningModule in primo piano

Incapsula l’architettura del modello
Organizza la logica di training in un’unica unità gestibile
Un blueprint che porta ordine e chiarezza ai progetti di deep learning

Diagramma di PyTorch LightningModule

Definire il metodo init

Task chiave:

Inizializzazione del modello
super():
- Gestione automatica dei training loop
- Logging
- Checkpointing
Layer del modello definiti dopo l’inizializzazione
Modulare e facile da mantenere

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)

Implementare il metodo forward

Passi chiave:

Definisci il flusso dei dati nella rete
Elabora l’input attraverso i layer in sequenza
- Trasformazione lineare
- Attivazione
- Ultimo layer e output

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

Esempio: classificare cifre scritte a mano

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)

Integrare il modello con task di classificazione

Focus sul caso d’uso di classificazione
Flusso completo dentro LightningModule
Output grezzi per l’attivazione softmax
Integrazione con 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

Ayo berlatih!

Modelli di AI scalabili con PyTorch Lightning