9.6 6. Pytorch lightning

딥러닝 실험을 구현하기 위해서는 뉴럴네트워크와 같은 모델 코드 외에도 그 시스템을 만들고 실험을 수행하기 위한 많은 엔지니어링 코드가 필요합니다. 이러한 코드들은 직접 짜는게 귀찮을뿐더러 남이 짠 코드를 읽을 때도 코드를 분석하기 어렵게 만듭니다. 그런데 사실 딥러닝에서의 많은 엔지니어링 코드는 모델이 달라져도 그 역할이 비슷비슷한 경우가 많습니다.

머신러닝, 딜러닝 모델 구축을 할 때에 1)공통된 부분들을 반복해서 작성할 필요 없이 대신 처리해주고, 2)머신러닝 모델 구축의 탬플릿 코드로써 기능을 하며, 3)다른 사람이 작성한 코드를 쉽게 볼 수 있도록 공통된 스타일을 갖도록 하고 4)모델의 개별적인 부분은 유연하게 커스터마이징하여 실험할 수 있는 라이브러리가 PyTorch Lightning 입니다.

PyTorch Lightning은 PyTorch에 대한 High-level 인터페이스를 제공하는 오픈소스 Python 라이브러리입니다. PyTorch만으로도 충분히 다양한 AI 모델들을 쉽게 생성할 수 있지만 GPU나 TPU, 그리고 16-bit precision, 분산학습 등 더욱 복잡한 조건에서 실험하게 될 경우, 코드가 복잡해집니다. 따라서 코드의 추상화를 통해, 프레임워크를 넘어 하나의 코드 스타일로 자리 잡기 위해 탄생한 프로젝트가 바로 PyTorch Lightning입니다.

기존 PyTorch는 DataLoader, Mode, optimizer, Training roof 등을 전부 따로따로 코드로 구현을 해야하는데 Pytorch Lightning에서는 Lightning Model class 안에 이 모든것을 한 번에 구현하도록 되어있습니다. 클래스 내부에 있는 함수명은 똑같이 써야하고 그 목적에 맞게 코딩해야 합니다.

Pytorch Lightning은 크게 Trainer와 Lightning Module로 나누어 살펴볼 수 있습니다.

Lightning Module은 모델 내부의 구조를 설계하는 research & science 클래스라고 생각할 수 있습니다. 모델의 구조나 데이터 전처리, 손실함수 등의 설정을 통해 모델을 초기화 합니다. 실제로 코드에서는 pl.LightningModule 클래스를 상속받아 새로운 LightningModule 클래스를 생성합니다. 기존 PyTorch의 nn.Module과 같은 방식이라고 보시면 됩니다.

Trainer는 모델의 학습을 담당하는 클래스라고 볼 수 있습니다. 모델의 학습 epoch이나 batch 등의 상태뿐만 아니라, 모델을 저장해 로그를 생성하는 부분까지 담당합니다. 실제로 코드에서는 pl.Trainer()라고 정의하면 끝입니다.

결국 두 가지의 큰 클래스를 통해, 복잡한 양의 작업들을 2가지 영역으로 추상화할 수 있게 됩니다.

1.1. PyTorch Lightning 간단 예제

PyTorch Lightning을 사용하여 딥러닝 모델을 작성하는 순서는

1) Lightning Module에서 상속된 새로운 Lightning Module 클래스를 작성합니다.

2) DataLoader 를 통해 학습할 데이터를 준비 합니다.

3) Trainer 객체를 만들고, 그 Trainer에 데이터와 Lightning Module 클래스를 주어 학습합니다.

import pytorch_lightning as pl

class LitModel(pl.LightningModule):
    def __init__(self):
        super().__init__()
        self.l1 = nn.Linear(28 * 28, 10)

    def forward(self, x):
        return torch.relu(self.l1(x.view(x.size(0), -1)))

    def training_step(self, batch, batch_idx):
        x, y = batch
        y_hat = self(x)
        loss = F.cross_entropy(y_hat, y)
        return loss

    def configure_optimizers(self):
        return torch.optim.Adam(self.parameters(), lr=0.02)

train_loader = DataLoader(MNIST(os.getcwd(), download=True, transform=transforms.ToTensor()))
trainer = pl.Trainer()
model = LitModel()
trainer.fit(model, train_loader)

1.2. 데이터 준비

PyTorch의 데이터 준비하는 과정을 크게 5가지 형태로 구조화하여 처리합니다.

1) 다운로드

2) 데이터 정리 혹은 메모리 저장

3) 데이터셋 로드

4) 데이터 전처리 (특히, transforms를 의미)

5) dataloader 형태로 wrapping

위의 순서에 맞게 코드를 추상화합니다.

prepare_data()

train_dataloader, val_dataloader, test_dataloader

PyTorch 학습 과정에 관여하는 여러 코드가 추상화된 함수형태로 LightningModule안에 포함되어 있습니다. 특히, 상속받은 LightningModule 클래스는 위와 같은 함수들을 순서에 따라 실행하는데, 이를 바로 Lifecycle이라고 부릅니다. 즉, 해당하는 순서에 따라 함수를 작성하는 것이 중요합니다.

1.3. LightningModule Class

Lightning Module은 5가지로 구성됩니다.

Computations (init).

Train loop (training_step)

Validation loop (validation_step)

Test loop (test_step)

Optimizers (configure_optimizers)

Lightning Module에는 많은 편리한 Method가 있지만 핵심적으로 알아야 하는 Method 들은 다음과 같습니다.

Name	Description
init	Define computations here
forward	Use for inference only (separate from training_step)
training_step	the full training loop
validation_step	the full validation loop
test_step	the full test loop
configure_optimizers	define optimizers and LR schedulers

1.3.1. training_step method

일반적인 pyTorch training loop 를 training_step 메서드내에 배치 합니다.

class LitClassifier(pl.LightningModule):
    def __init__(self, model):
        super().__init__()
        self.model = model

    def training_step(self, batch, batch_idx):
        x, y = batch
        y_hat = self.model(x)
        loss = F.cross_entropy(y_hat, y)
        return loss

만약 epoch-level metric을 계산하고 log를 하려면 .log 메서드를 사용합니다.

def training_step(self, batch, batch_idx):
    x, y = batch
    y_hat = self.model(x)
    loss = F.cross_entropy(y_hat, y)

    # logs metrics for each training_step,
    # and the average across the epoch, to the progress bar and logger
    self.log("train_loss", loss, on_step=True, on_epoch=True, prog_bar=True, logger=True)
    return loss

만약에 각 training_step의 결과로 무엇인가 할 일이 있으면 training_epoch_end 메서드에 작성합니다.

def training_step(self, batch, batch_idx):
    x, y = batch
    y_hat = self.model(x)
    loss = F.cross_entropy(y_hat, y)
    preds = ...
    return {"loss": loss, "other_stuff": preds}


def training_epoch_end(self, training_step_outputs):
    for pred in training_step_outputs:
        ...

validation loop를 추가 하려면 validation_step 메서드를 추가 합니다.

class LitModel(pl.LightningModule):
    def validation_step(self, batch, batch_idx):
        x, y = batch
        y_hat = self.model(x)
        loss = F.cross_entropy(y_hat, y)
        self.log("val_loss", loss)

만약에 각 validation _step의 결과로 무엇인가 할 일이 있으면 validation _epoch_end 메서드에 작성합니다.

def validation_step(self, batch, batch_idx):
    x, y = batch
    y_hat = self.model(x)
    loss = F.cross_entropy(y_hat, y)
    pred = ...
    return pred


def validation_epoch_end(self, validation_step_outputs):
    for pred in validation_step_outputs:

1.4. MNIST 예제

이제 PyTorch Lightning 을 사용한 MNIST 학습 코드를 설명해 보겠습니다.

먼저 필요한 모듈 몇 가지를 설치해야 합니다.

!pip install --quiet torchmetrics torch “pytorch-lightning>=1.3" "torchvision"

작성되는 코드에서 사용할 라이브러리들을 import 합니다.

import os

import torch
import pytorch_lightning as pl
from pytorch_lightning import LightningModule, Trainer
from torch import nn
from torch.nn import functional as F
from torch.utils.data import DataLoader, random_split
from torchvision import transforms
from torchvision.datasets import MNIST

PATH_DATASETS = os.environ.get("PATH_DATASETS", ".")
AVAIL_GPUS = min(1, torch.cuda.device_count())
BATCH_SIZE = 256 if AVAIL_GPUS else 64

다음은 최소화된 MNIST 코드의 예제 입니다. training_step과 configure_optimizers 메서드는 필수적으로 구현해야 합니다.

class MNISTModel(LightningModule):
    def __init__(self):
        super().__init__()
        self.model = nn.Sequential(
            nn.Flatten(),
            nn.Linear(28 * 28, 64),
            nn.BatchNorm1d(64),
            nn.ReLU(inplace=True),
            nn.Linear(64, 64),
            nn.BatchNorm1d(64),
            nn.ReLU(inplace=True),
            nn.Linear(64, 10)
        )

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

    def training_step(self, batch, batch_nb):
        x, y = batch
        logits = self.model(x)
        loss = F.cross_entropy(logits, y)
        return loss

    def validation_step(self, batch, batch_idx):
        x, y = batch
        logits = self.model(x)
        loss = F.cross_entropy(logits, y)
        self.log("train_loss", loss)
        return loss

    def test_step(self, batch, batch_idx):
        x, y = batch
        logits = self.model(x)
        loss = F.cross_entropy(logits, y)
        self.log("test_loss", loss)
        return loss

    def configure_optimizers(self):
        return torch.optim.Adam(self.parameters(), lr=0.02)

__init__ 에서 모델의 구조를 정의합니다. MNIST 를 위해 몇 개의 층을 Sequential로 쌓았습니다.

l forward는 모델의 추론 결과를 제공하고 싶을 때 사용합니다. nn.Module처럼 꼭 정의해야 하는 메서드는 아니지만 self(<입력>)과 같이 사용할 수 있게 만들어주므로 구현해주면 다른 메서드를 구현할 때 편리합니다.

l training_step은 학습 루프의 body 부분을 나타냅니다. 이 메소드에서는 argument로 training 데이터로더가 제공하는 batch와 해당 batch의 인덱스가 주어지고 학습 로스를 계산하여 리턴합니다. pytorch lightning은 편리하게도 batch의 텐서를 cpu 혹은 gpu 텐서로 변경하는 코드를 따로 추가하지 않아도 trainer의 설정에 따라 자동으로 적절한 타입으로 변경해줍니다.

l validation_step은 학습 중간에 모델의 성능을 체크하는 용도로 사용합니다. training_step과 마찬가지로 validation 데이터로더에서 제공하는 배치를 가지고 확인하고자 하는 통계량을 기록할 수 있습니다. 하나의 값을 저장할 때는 self.log(<변수 이름="">, <값>)과 같이 저장할 수 있고 여러 개의 변수를 저장하고 싶으면 아래 예시와 같이 self.log_dict로 변수 이름, 값 쌍을 가지고 있는 딕셔너리를 저장할 수 있습니다. 각 스탭마다 변수에 저장된 값의 평균이 해당 변수의 최종 값이 됩니다. 특별히 설정을 바꾸지 않으면 변수 중에 'val_loss'가 best 모델을 구하는 기준으로 사용됩니다. 예제에서는 모델의 정확도와 cross entropy loss를 구해서 저장합니다. 여기서 accuracy 함수는 pytorch_lightning.metrics.functional에서 정의되어 있는 함수로 logits에서 최댓값인 라벨이 실제 라벨과 일치하는 비율을 구해줍니다.

l configure_optimizers에서는 모델의 최적 파라미터를 찾을 때 사용할 optimizer와 scheduler를 구현합니다. GANs와 같이 여러 모델을 학습하기 위해 여러 optimizer를 사용해야 한다면 리스트로 리턴하면 됩니다. 이 경우에는 training_step에서 optimizer의 인덱스를 추가로 받아서 여러 모델을 번갈아 학습하게 됩니다. 예제에서는 학습해야 할 모델이 하나이므로 하나의 Adam optimzer만 사용하도록 하겠습니다.

pl.seed_everything(777)

# Init DataLoader from MNIST Dataset
dataset = MNIST(PATH_DATASETS, train=True, download=True, transform=transforms.ToTensor())

train_dataset, val_dataset = random_split(dataset, [55000, 5000])
test_dataset = MNIST('', train=False, download=True, transform=transforms.ToTensor())
train_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE)
val_loader = DataLoader(val_dataset, batch_size=BATCH_SIZE)
test_loader = DataLoader(test_dataset, batch_size=BATCH_SIZE)

# Initialize a trainer
trainer = Trainer(gpus=AVAIL_GPUS, max_epochs=3, progress_bar_refresh_rate=20,)

# Init our model
mnist_model = MNISTModel()

# Train the model ⚡
trainer.fit(mnist_model, train_loader)

trainer.test(test_dataloaders=test_loader)

학습 및 테스트에 사용할 MNIST 데이터셋을 불러옵니다. 여기서 학습 데이터셋의 일부를 랜덤으로 샘플링해 validation 용도로 사용하겠습니다. 각 나눈 데이터 셋을 가지고 training, validation, test용 데이터로더를 만듭니다.

모델을 학습하기 위해서는 학습 로직을 정하는 Trainer를 생성해야 합니다. Pytorch lightning의 Trainer는 굉장히 많은 기능을 제공합니다. 아래 예제에서는 간단히 학습 epoch 수와 gpu 수만 조정할 수 있도록 만들었습니다.

gpus가 0일 때는 cpu를 사용하고 gpus가 1 이상이면 gpu를 사용하여 모델을 학습합니다. gpus가 2 이상이면 자동으로 다중 gpu를 활용해 분산 학습을 진행하게 되는데 기본 설정은 process를 spawn하는 distributed data parallel 방식(ddp_spawn)으로 되어있습니다.

Trainer와 lightning module을 정의하고 난 뒤에 Trainer의 fit 함수로 모델을 학습할 수 있습니다. fit의 파라미터로 모델, training 데이터로더와 validation 데이터로더를 넘겨줍니다.

Pytorch lightning은 기본적으로 각 버전마다 체크포인트를 저장해줍니다. 하지만, 체크포인트 이름, 저장 주기, 모니터링할 metric 등을 바꾸고 싶으면 체크포인트 callback을 수정해주어야 합니다. 아래와 같이 ModelCheckpoint로 체크포인트 콜백을 생성합니다.

checkpoint_callback = ModelCheckpoint(

filepath=os.path.join('checkpoints', '{epoch:d}'),

verbose=True,

save_last=True,

save_top_k=args.save_top_k,

monitor='val_acc',

mode='max'

)

각 파라미터의 의미는 다음과 같습니다.

- filepath: 체크포인트 저장위치와 이름 형식을 지정합니다.

- verbose: 체크포인트 저장 결과를 출력합니다.

- save_last: 마지막 체크포인트를 저장합니다.

- save_top_k: 최대 몇 개의 체크포인트를 저장할지 지정합니다.(save_last에 의해 저장되는 체크포인트는 제외)

- monitor: 어떤 metric을 기준으로 체크포인트를 저장할지 지정합니다.

- mode: 지정한 metric의 어떤 기준(ex. min, max)으로 체크포인트를 저장할지 지정합니다.

Trainer 에 CheckPoint 를 적용하면 다음과 같이 코드를 수정해야 합니다.

from pytorch_lightning.callbacks import ModelCheckpoint

pl.seed_everything(777)

# Init DataLoader from MNIST Dataset
dataset = MNIST(PATH_DATASETS, train=True, download=True, transform=transforms.ToTensor())

train_dataset, val_dataset = random_split(dataset, [55000, 5000])
test_dataset = MNIST('', train=False, download=True, transform=transforms.ToTensor())
train_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE)
val_loader = DataLoader(val_dataset, batch_size=BATCH_SIZE)
test_loader = DataLoader(test_dataset, batch_size=BATCH_SIZE)

checkpoint_callback = ModelCheckpoint(
    filename='sample-mnist-epoch{epoch:02d}-val_loss{val/loss:.2f}',
    verbose=True,
    save_last=True,
    monitor='val_acc',
    mode='max'
)

# Initialize a trainer
trainer = Trainer(gpus=AVAIL_GPUS, max_epochs=3, progress_bar_refresh_rate=20,callbacks=[checkpoint_callback])

# Init our model
mnist_model = MNISTModel()

# Train the model ⚡
trainer.fit(mnist_model, train_loader)