[Pytorch Lightning] 튜토리얼 1 - 모델을 만들고 훈련시키기

오늘은 파이토치 라이트닝과 관련해 기본적인 자료가 부족해서 직접 작성해보기로 했다. 공식 홈페이지를 기반으로 했으니, 많이 도움이 될 것이라고 생각한다. 코드는 여기를 참고하면 된다.

Module, Trainer는 Pytorch Lightning이 보라색이라서 그렇게 색깔을 정했다

대략적인 구조는 다음과 같다. 

1. 데이터를 모아서 DataLoader를 이용해 train, valid, test로 분할한다.
2. Model을 선언한다.
3. 훈련때마다 사용하는 방법을 선언한다. (e.g., train_step에서는 어떻게 진행하고 ...)
4. Model과 훈련 방법을 Pytorch Lightning Module에 선언한다.
5. Hyper Parameter + Lighitning Module을 Trainer에 추가한다.
6. Trainer를 통해 훈련된 모델(Trained Model)을 사용한다.

 

요약

• 라이브러리 설치
• 데이터셋 선언
• 모델 구조 만들기
• Lighitning Module 만들기
• Trainer를 이용해 훈련 및 테스트하기

5가지 단계로 구성되어 있다.

 

라이브러리 설치 및 선언

pip install torch==2.3.1
pip install torchvision==0.18.1
pip install pytorch_lightning==2.3.3

 

# 필요한 라이브러리 선언

import os
import torch
import numpy as np
import matplotlib.pyplot as plt
from torch import nn
import torch.nn.functional as F
from torchvision import transforms
from torchvision.datasets import MNIST
from torch.utils import data
import pytorch_lightning as pl

 

데이터셋 선언(분할) 및 확인

dataset은 손글씨로 유명헌 MNIST 데이터셋을 사용한다. 자세한 정보는 여기를 확인하면 된다.

데이터셋 선언 및 분할

# 데이터 셋 불러오기
train_set = MNIST(os.getcwd(), download=True, train=True, transform=transforms.ToTensor())
test_set = MNIST(os.getcwd(), download=True, train=False, transform=transforms.ToTensor())

# 데이터셋 분할
train_set_size = int(len(train_set) * 0.8)
valid_set_size = len(train_set) - train_set_size

# 8 : 2 비율로 train_set과 valid_set 분할
train_set, valid_set = data.random_split(train_set, [train_set_size, valid_set_size])

데이터 확인

image, label = train_set[0]
image = image.numpy()
image = image.squeeze()  # 차원 축소

# 이미지 시각화
plt.imshow(image, cmap='gray')
plt.title(f"answer = {label}")
plt.show()

데이터로더 변환

# DataLoader 변환
train_loader = data.DataLoader(train_set)
valid_loader = data.DataLoader(valid_set)
test_loader = data.DataLoader(test_set)

 

모델 만들기(Encoder, Decoder)

class Encoder(nn.Module):
    def __init__(self):
        super().__init__()
        self.l1 = nn.Sequential(nn.Linear(28 * 28, 64), nn.ReLU(), nn.Linear(64, 3))

    def forward(self, x):
        return self.l1(x)


class Decoder(nn.Module):
    def __init__(self):
        super().__init__()
        self.l1 = nn.Sequential(nn.Linear(3, 64), nn.ReLU(), nn.Linear(64, 28 * 28))

    def forward(self, x):
        return self.l1(x)

Encoder 

• 목적: 입력 데이터를 저차원의 표현으로 변환(인코딩).
• 구조:
  • nn.Linear(28 * 28, 64): 입력 데이터의 차원을 784(28x28)에서 64로 줄이는 선형 변환을 수행한다.
  • nn.ReLU(): 비선형 활성화 함수로, 선형 변환 후 적용되어 모델의 표현력을 증가시킨다.
  • nn.Linear(64, 3): 중간 표현의 차원을 64에서 3으로 더 줄여 저차원 특징을 추출한다.

Decoder 클래스

• 목적: 저차원 표현을 다시 원래의 차원으로 확장하여 데이터를 복원(디코딩).
• 구조:
  • nn.Linear(3, 64): 저차원 표현(3차원)을 중간 차원(64차원)으로 확장하는 선형 변환을 수행한다.
  • nn.ReLU(): 비선형 활성화 함수를 적용하여 데이터의 비선형적 특징을 유지 및 강화한다.
  • nn.Linear(64, 28 * 28): 최종적으로 중간 차원을 원래의 입력 데이터 차원(784차원)으로 확장하여 재구성한다.

 

Lighitning Module 만들기

class LitAutoEncoder(pl.LightningModule):
    def __init__(self, encoder, decoder):
        super().__init__()
        self.encoder = encoder
        self.decoder = decoder

    def training_step(self, batch, batch_idx):
        # training_step defines the train loop.
        x, _ = batch
        x = x.view(x.size(0), -1)
        z = self.encoder(x)
        x_hat = self.decoder(z)
        loss = F.mse_loss(x_hat, x)
        return loss

    def validation_step(self, batch, batch_idx):
        # this is the validation loop
        x, _ = batch
        x = x.view(x.size(0), -1)
        z = self.encoder(x)
        x_hat = self.decoder(z)
        val_loss = F.mse_loss(x_hat, x)
        self.log("val_loss", val_loss)
    
    def test_step(self, batch, batch_idx):
        # this is the test loop
        x, _ = batch
        x = x.view(x.size(0), -1)
        z = self.encoder(x)
        x_hat = self.decoder(z)
        test_loss = F.mse_loss(x_hat, x)
        self.log("test_loss", test_loss)

    def configure_optimizers(self):
        optimizer = torch.optim.Adam(self.parameters(), lr=1e-3)
        return optimizer

각각 함수에 대한 설명과 코드 해석을 하면 다음과 같다.

• __init__(self, encoder, decoder): 생성자에서는 인코더와 디코더 객체를 모듈의 속성으로 저장한다.
• training_step : 훈련 step마다 실행되며, x 입력에 대해 인코딩, 디코딩과정을 통해 정답과 예측에 대해 loss 계산한다.
• validation_step : 검증 step마다 실행되며, training_step과 동일하다.
• test_step : 테스트 step에서 실행되며, training_step과 동일하다.
• configure_optimizers : 학습 과정에서 사용할 옵티마이저를 설정한다. 필요하다면 학습률 스케줄러도 설정할 수 있지만, 여기서는 0.001로 설정한다.

 

모듈을 이용한 모델 선언

# model
model = LitAutoEncoder(Encoder(), Decoder())

 

Trainer로 훈련 및 테스트 하기

 

Trainer 선언

trainer = pl.Trainer(max_epochs=5)

Trainer 내부에 Hyper Parameters 선언하지만, 지금은 간단하게 epoch만 설정했다.

Trainer로  모델 훈련하기

trainer.fit(model, train_loader, valid_loader)

trainer의 fit 함수를 이용하여 train_loader와 valid_loader를 이용해 model 훈련한다.

• train_loader : training_step에서 사용된다.
• valid_loader : validation_step에서 사용된다.

실행하면 다음과 같이 모델 정보(파라미터 수, GPU 정보) 와 함께 실험 예상시간이 출력된다.

 

Trainer로 모델 테스트하기

trainer.test(model, dataloaders=test_loader)

실행하면 실행에 대한 결과가 출력된다. 여기서 코드 작성을 통해 원하는 metric을 추가할 수 있다.

 

다음은 훈련된 모델을 불러오고 사용하는 방법에 대해 알아보려고 한다.