어느덧 벌써 12번째 인먹 스터디!!
오늘의 주제는 바로!!! Pytorch Lightnig이다.
그럼 같이 가보자고!
https://lightning.ai/docs/pytorch/stable/
딥러닝을 구현할 수 있는 다양한 프레임워크가 나오고 있다. 파이썬에서 우리는 대표적으로 tensorflow, pytorch를 이용하여 딥러닝을 구현한다. 딥러닝을 다루다보면 꼭 필요한 과정이 존재하는데 예를 들어 model 구현, loss, optimizer 정의, dataset & dataloader 정의, epoch별 학습 등이 있다.
매 구현에서 동일한 과정이 꼭 따라오는데 소개할 Pytorch Lightning은 딥러닝 구축 시 공통된 부분을 반복해서 사용할 필요가 없고, 다른 사람이 작성한 코드를 쉽게 볼 수 있도록 공통된 스타일의 모듈로 구현된다.
다양한 장점이 존재하는데 아래에서 코드와 함께 살펴보기로 한다. 굳이 한계점이라 하면, class를 잘 정의할 줄 아는 사람이어야 할 것 같다. 이번 스터디를 통해 공부해야 할 필요성을 느꼈다.
그럼 시작해볼까요 ~?
먼저, Pytorch Lightning을 설치하는 방법은 다음과 같다.
pip 설치
'''pip install lightning'''
conda 설치
'''conda install lightning -c conda-forge'''
둘 중 편한걸로 깔아주면 되는데, 작성 시점에서의 Pytorch Lightning 버전은 2.2.0이다.
들어가기 전, 작성자는 파이토치를 다룬지 꽤 되었다!
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)
공식문서에 있는 예시를 가져왔다. 여기에 존재하는 여러 정의된 함수를 살펴보면 다음과 같다.
2. Trainer 정의
Lightning Trainer는 모든 LightningModule와 혼합하고 확장에 필요한 모든 엔지니어링의 복잡성을 추상화할 수 있다. Trainer 안에는 다양한 옵션들이 존재하고 있는데 알아보자.
import lightning as L
trainer = L.Trainer()
trainer.fit(model=autoencoder, train_dataloaders=train_loader)
Pytorch Lightning의 가장 큰 장점 중 하나는 자동으로 가중치를 저장해주고 텐서보드를 이용하여 학습 그래프를 시각화 가능 하다는 점이다. 기존에서는 가중치 저장 코드, 텐서보드 저장 코드를 다 정의해주어야 했다면 자동 활성화가 가능하다.
checkpoint = "./lightning_logs/version_0/checkpoints/epoch=0-step=100.ckpt"
autoencoder = LitAutoEncoder.load_from_checkpoint(checkpoint, encoder=encoder, decoder=decoder)
def custom_callbacks(x):
return x
L.Trainer(callbacks=[custom_callbacks])
정의하여 옵션으로 넣어주면 된다. 위에 관해서는 바로 뒤에 나온다.
반대로, 저장하고 싶지 않다면? 끄는 기능도 옵션에서 찾으면 된다.
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__에 추가하면 된다.
일부를 저장하고 싶다면 하이퍼파라미터를 옵션처럼 적어주면 되고 그게 아니라면 저 옵션만 써줘도 전체가 저장된다고 한다.
아디오스!!!😊✨