Định nghĩa mô hình với LightningModule

Mô hình AI có khả năng mở rộng với PyTorch Lightning

Sergiy Tkachuk

Director, GenAI Productivity

LightningModule trọng tâm

Đóng gói kiến trúc mô hình
Tổ chức logic huấn luyện trong một khối dễ quản lý
Bản thiết kế giúp dự án deep learning rõ ràng, có trật tự

Sơ đồ PyTorch LightningModule

Định nghĩa phương thức init

Nhiệm vụ chính:

Khởi tạo mô hình
super():
- Xử lý vòng lặp huấn luyện tự động
- Ghi log
- Lưu checkpoint
Định nghĩa các tầng sau khi khởi tạo
Mô-đun hóa, dễ bảo trì

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)

Cài đặt phương thức forward

Các bước chính:

Xác định luồng dữ liệu qua mạng
Xử lý đầu vào tuần tự qua các tầng
- Biến đổi tuyến tính
- Kích hoạt
- Tầng cuối và đầu ra

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

Ví dụ: phân loại chữ số viết tay

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)

Tích hợp mô hình cho tác vụ phân loại

Tập trung vào bài toán phân loại
Toàn bộ luồng trong LightningModule
Trả về đầu ra thô cho softmax
Tích hợp với 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!

Mô hình AI có khả năng mở rộng với PyTorch Lightning