Modellen definiëren met LightningModule

Schaalbare AI-modellen met PyTorch Lightning

Sergiy Tkachuk

Director, GenAI Productivity

LightningModule uitgelicht

Vat je modelarchitectuur samen
Bundelt trainingslogica in één beheersbare unit
Blauwdruk die orde en duidelijkheid brengt in deep‑learningprojecten

PyTorch LightningModule-diagram

De init-methode definiëren

Belangrijkste taken:

Initialisatie van het model
super():
- Automatisch trainingsloops afhandelen
- Logging
- Checkpointing
Lagen definiëren na initialisatie
Modulair en onderhoudbaar

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)

De forward-methode implementeren

Stappen:

Definieer de gegevensstroom door het netwerk
Verwerk input sequentieel door de lagen
- Lineaire transformatie
- Activatie
- Laatste laag en 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

Voorbeeld: handgeschreven cijfers classificeren

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)

Het model integreren met classificatietaken

Focus op classificatie-usecase
Volledige flow binnen LightningModule
Geeft ruwe outputs voor softmax-activatie
Integratie met 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

Laten we oefenen!

Schaalbare AI-modellen met PyTorch Lightning