Mendefinisikan model dengan LightningModule

Model AI yang Dapat Diskalakan dengan PyTorch Lightning

Sergiy Tkachuk

Director, GenAI Productivity

Fokus pada LightningModule

Mengenkapsulasi arsitektur model Anda
Mengatur logika pelatihan dalam satu unit yang mudah dikelola
Cetak biru yang memberi keteraturan dan kejelasan pada proyek deep learning

Diagram PyTorch LightningModule

Mendefinisikan metode init

Tugas utama:

Inisialisasi model
super():
- Penanganan loop pelatihan otomatis
- Logging
- Checkpointing
Layer model didefinisikan setelah inisialisasi
Modular dan mudah dirawat

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)

Mengimplementasikan metode forward

Langkah kunci:

Definisikan alur data melalui jaringan
Proses input berurutan melalui layer
- Transformasi linear
- Aktivasi
- Layer terakhir dan 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

Contoh: klasifikasi digit tulisan tangan

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)

Mengintegrasikan model untuk tugas klasifikasi

Fokus pada kasus klasifikasi
Alur lengkap dalam LightningModule
Keluarkan logit untuk aktivasi softmax
Integrasi dengan 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!

Model AI yang Dapat Diskalakan dengan PyTorch Lightning