[인먹] Pytorch Lightning (12차)

jihye_heo·2024년 2월 28일
0

Study

목록 보기
4/4

어느덧 벌써 12번째 인먹 스터디!!
오늘의 주제는 바로!!! Pytorch Lightnig이다.
그럼 같이 가보자고!


Pytorch Lightning

https://lightning.ai/docs/pytorch/stable/

딥러닝을 구현할 수 있는 다양한 프레임워크가 나오고 있다. 파이썬에서 우리는 대표적으로 tensorflow, pytorch를 이용하여 딥러닝을 구현한다. 딥러닝을 다루다보면 꼭 필요한 과정이 존재하는데 예를 들어 model 구현, loss, optimizer 정의, dataset & dataloader 정의, epoch별 학습 등이 있다.

매 구현에서 동일한 과정이 꼭 따라오는데 소개할 Pytorch Lightning은 딥러닝 구축 시 공통된 부분을 반복해서 사용할 필요가 없고, 다른 사람이 작성한 코드를 쉽게 볼 수 있도록 공통된 스타일의 모듈로 구현된다.

다양한 장점이 존재하는데 아래에서 코드와 함께 살펴보기로 한다. 굳이 한계점이라 하면, class를 잘 정의할 줄 아는 사람이어야 할 것 같다. 이번 스터디를 통해 공부해야 할 필요성을 느꼈다.

그럼 시작해볼까요 ~?

Install

먼저, Pytorch Lightning을 설치하는 방법은 다음과 같다.

  • pip 설치
    '''pip install lightning'''

  • conda 설치
    '''conda install lightning -c conda-forge'''

둘 중 편한걸로 깔아주면 되는데, 작성 시점에서의 Pytorch Lightning 버전은 2.2.0이다.

들어가기 전, 작성자는 파이토치를 다룬지 꽤 되었다!

Tutorial

Pytorch Lightning의 공식 문서에서는 다양한 예제 튜토리얼이 존재한다. 오늘은 이를 따라해보며 Lightning과 친해지는 시간을 가져보려한다.

링크 : https://lightning.ai/docs/pytorch/stable/starter/introduction.html

Pytorch lightning은 크게 LightningModule와 Trainer로 나눠서 살펴볼 수 있다.

1. LightningModule 정의

LightningModule은 정의된 모델과 함께 학습, 검증, 테스트 단계의 step을 함께 정의해주는 모듈이다. 기존 pytorch의 nn.Module과 같은 방식으로 정의된다.

import os
from torch import optim, nn, utils, Tensor
from torchvision.datasets import MNIST
from torchvision.transforms import ToTensor
import lightning as L

# define any number of nn.Modules (or use your current ones)
encoder = nn.Sequential(nn.Linear(28 * 28, 64), nn.ReLU(), nn.Linear(64, 3))
decoder = nn.Sequential(nn.Linear(3, 64), nn.ReLU(), nn.Linear(64, 28 * 28))


# define the LightningModule
class LitAutoEncoder(L.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.
        # it is independent of forward
        x, y = batch
        x = x.view(x.size(0), -1)
        z = self.encoder(x)
        x_hat = self.decoder(z)
        loss = nn.functional.mse_loss(x_hat, x)
        # Logging to TensorBoard (if installed) by default
        self.log("train_loss", loss)
        return loss

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


# init the autoencoder
autoencoder = LitAutoEncoder(encoder, decoder)

공식문서에 있는 예시를 가져왔다. 여기에 존재하는 여러 정의된 함수를 살펴보면 다음과 같다.

  • training_step : dataloader를 for문으로 만드는 학습 과정을 정의한 부분, 당연히 validation_step과 test_step도 만들 수 있음.
  • configure_optimizers : optimizer를 정의하는 부분

2. Trainer 정의

Lightning Trainer는 모든 LightningModule와 혼합하고 확장에 필요한 모든 엔지니어링의 복잡성을 추상화할 수 있다. Trainer 안에는 다양한 옵션들이 존재하고 있는데 알아보자.

import lightning as L
trainer = L.Trainer()
trainer.fit(model=autoencoder, train_dataloaders=train_loader)

Pytorch Lightning의 가장 큰 장점 중 하나는 자동으로 가중치를 저장해주고 텐서보드를 이용하여 학습 그래프를 시각화 가능 하다는 점이다. 기존에서는 가중치 저장 코드, 텐서보드 저장 코드를 다 정의해주어야 했다면 자동 활성화가 가능하다.

  • 학습된 가중치를 load 하는 방법
    우선 학습을 시도하게 되면, 시도 횟수에 따른 버전 경로에 저장되게 된다.
    이후 load_from_checkpoint를 이용하여 가중치를 load할 수 있다.
checkpoint = "./lightning_logs/version_0/checkpoints/epoch=0-step=100.ckpt"
autoencoder = LitAutoEncoder.load_from_checkpoint(checkpoint, encoder=encoder, decoder=decoder)

  • 내장된 가중치 저장 방법말고 나만의 callback 함수로 정의하고 싶다면?
def custom_callbacks(x):
	return x
L.Trainer(callbacks=[custom_callbacks])

정의하여 옵션으로 넣어주면 된다. 위에 관해서는 바로 뒤에 나온다.
반대로, 저장하고 싶지 않다면? 끄는 기능도 옵션에서 찾으면 된다.

Level Up

https://lightning.ai/docs/pytorch/stable/expertise_levels.html

다음은 base에서 벗어난 level up 페이지.

참고로 Pytorch Lightning은 pytorch와 호환이 가능하다는 엄청 큰 장점이 함께 존재한다.

1) validation_step, test_step

validation_step을 train_step과 함께 쓰고 싶다면 fit에 같이 넣으면 된다.

from torch.utils.data import DataLoader

train_loader = DataLoader(train_set)
valid_loader = DataLoader(valid_set)
model = LitAutoEncoder(...)

# train with both splits
trainer = L.Trainer()
trainer.fit(model, train_loader, valid_loader)

test_step의 경우 fit이 아닌 test를 이용하여 결과를 도출한다.

# initialize the Trainer
trainer = Trainer()

# test the model
trainer.test(model, dataloaders=DataLoader(test_set))

2) checkpoint 경로 변경
앞에서 버전 횟수에 따라 저장되는 경로가 다르다고 하였다.
저장되는 root 경로를 바꿀 수 있는데 Trainer 옵션에서 가능하다.

# saves checkpoints to 'some/path/' at every epoch end
trainer = Trainer(default_root_dir="some/path/")

3) early stopping
기존 정의된 early stopping은 validation metric을 보고 성능이 오르지 않을 때 중지하게 되어 있다.

from pytorch_lightning.callbacks.early_stopping import EarlyStopping


def validation_step(self):
    self.log("val_loss", loss)


trainer = Trainer(callbacks=[EarlyStopping(monitor="val_loss")])

4) argumentparser

from argparse import ArgumentParser

parser = ArgumentParser()
parser.add_argument("--layer_1_dim", type=int, default=128)
args = parser.parse_args()

이렇게 정의했으면 나중에 py 파일을 실행할 때 값을 실행시킬 수 있다.

python trainer.py --layer_1_dim 64

그리고 다양한 argument로 받아올 수 있다.

5) hyperparameter 저장

이 기능까지 있어서 사실 난 되게 감탄했다. 방법도 매우 단순하다.
self.save_hyperparameters를 LightningModule의 __init__에 추가하면 된다.
일부를 저장하고 싶다면 하이퍼파라미터를 옵션처럼 적어주면 되고 그게 아니라면 저 옵션만 써줘도 전체가 저장된다고 한다.




이렇게 간단하게 알아보았다!! 다음 시간에는 Pytorch Lightning을 응용하여 데이터에 적용해 결과를 가져오려고 한다. 뭔가 파이토치랑 비슷한데 다른 부분이 많아서 재밌게 공부를 할 수 있었던 것 같다.

아디오스!!!😊✨

0개의 댓글

관련 채용 정보