LightningModule로 모델 정의하기

PyTorch Lightning으로 만드는 확장 가능한 AI 모델

Sergiy Tkachuk

Director, GenAI Productivity

LightningModule 핵심

  1. 모델 아키텍처를 캡슐화
  2. 학습 로직을 단일, 관리 가능한 단위로 구성
  3. 딥러닝 프로젝트에 질서와 명확성을 주는 청사진

PyTorch LightningModule 다이어그램

PyTorch Lightning으로 만드는 확장 가능한 AI 모델

__init__ 메서드 정의하기

핵심 작업:

  • 모델 초기화
  • super():
    • 학습 루프 자동 처리
    • 로깅
    • 체크포인트
  • 초기화 후 모델 층 정의
  • 모듈식이며 유지보수 용이
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)
PyTorch Lightning으로 만드는 확장 가능한 AI 모델

forward 메서드 구현하기

핵심 단계:

  • 네트워크의 데이터 흐름 정의
  • 입력을 순차적으로 처리
    • 선형 변환
    • 활성화
    • 마지막 층과 출력
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
PyTorch Lightning으로 만드는 확장 가능한 AI 모델

예시: 손글씨 숫자 분류

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)
PyTorch Lightning으로 만드는 확장 가능한 AI 모델

모델을 분류 작업에 통합하기

$$

  • 분류(use case)에 집중
  • 전체 흐름을 LightningModule에 구현
  • 소프트맥스 전에 원시 출력 반환
  • 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
PyTorch Lightning으로 만드는 확장 가능한 AI 모델

연습해 봅시다!

PyTorch Lightning으로 만드는 확장 가능한 AI 모델

Preparing Video For Download...