การกำหนดโมเดลด้วย LightningModule

AI Models แบบ Scalable ด้วย PyTorch Lightning

Sergiy Tkachuk

Director, GenAI Productivity

มองโฟกัสที่ LightningModule

  1. ครอบคลุมสถาปัตยกรรมโมเดล
  2. จัดระเบียบตรรกะการฝึกในหน่วยเดียวที่จัดการได้ง่าย
  3. พิมพ์เขียวที่นำความเป็นระเบียบและความชัดเจนมาสู่โปรเจกต์ deep learning

แผนภาพ PyTorch LightningModule

AI Models แบบ Scalable ด้วย PyTorch Lightning

การกำหนด method `__init__`

งานหลัก:

  • การกำหนดค่าเริ่มต้นของโมเดล
  • super():
    • จัดการ training loop อัตโนมัติ
    • การบันทึก log
    • การบันทึก checkpoint
  • กำหนด layer ของโมเดลหลังการ initialize
  • เป็นโมดูลและดูแลรักษาง่าย
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)
AI Models แบบ Scalable ด้วย PyTorch Lightning

การสร้าง method `forward`

ขั้นตอนหลัก:

  • กำหนดการไหลของข้อมูลผ่านเครือข่าย
  • ประมวลผล input ผ่านแต่ละ layer ตามลำดับ
    • การแปลงเชิงเส้น
    • Activation
    • Layer สุดท้ายและ 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
AI Models แบบ Scalable ด้วย PyTorch Lightning

ตัวอย่าง: จำแนกตัวเลขลายมือ

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)
AI Models แบบ Scalable ด้วย PyTorch Lightning

การนำโมเดลไปใช้กับงานจำแนกประเภท

$$

  • เน้น use case การจำแนกประเภท
  • ขั้นตอนทั้งหมดอยู่ใน LightningModule
  • ส่งออก raw output สำหรับ softmax activation
  • ทำงานร่วมกับ 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
AI Models แบบ Scalable ด้วย PyTorch Lightning

มาฝึกกันเถอะ!

AI Models แบบ Scalable ด้วย PyTorch Lightning

Preparing Video For Download...