PyTorch 쓰세요.
#pytorch #pytorch_lightning
파이토치 쓰세요.
당장 쓰세요.
하던거 멈추고 이거 먼저 쓰세요. 어서.
PyTorch Lightning 관련해서 모델을 만드는데,
너무나 편리(?)하게 코드를 작성하고 지극히 개인적인 생각으로는
이젠 누구나 모델을 만들고 학습시키면서
각각 자신만의 AI 환경을 만들 수 있겠다는 생각이 문득 떠올랐다.
What is the PyTorch Lightning?
https://pytorch-lightning.readthedocs.io/en/latest/
Doc 을 먼저 다~ 읽어보면 나와 같은 생각을 할 수 있을 것 같다.
딥러닝 분야에서 Hello World! 에 해당하는 Mnist 손글씨 숫자 맞추기 게임을 통해서
PyTorch Lightning 에 대해서 간단하게 살펴보면 좋겠다.
(사이트에서 제공하는 예제이다.)
나의 환경은,
GPU 2장 - 1 Node (컴퓨터 한대)
쥬피터 노트북(Python 3.6)
먼저 import 불롸불롸~
import os
import torch
from pytorch_lightning import LightningModule, Trainer
from pytorch_lightning.metrics.functional import accuracy
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', './data/')
# 저장경로는 필요한 장소로 변경 후 사용
AVAIL_GPUS = max(1, torch.cuda.device_count())
# 예제에는 min 으로 되어있어서 수정함
BATCH_SIZE = 256 if AVAIL_GPUS else 64
여기까진 쉽다.
설명도 필요없다.
바로 모델을 만드는데,
class LitMNIST(LightningModule):
def __init__(self, data_dir=PATH_DATASETS, hidden_size=64, learning_rate=2e-4):
super().__init__()
# Set our init args as class attributes
self.data_dir = data_dir
self.hidden_size = hidden_size
self.learning_rate = learning_rate
# Hardcode some dataset specific attributes
self.num_classes = 10
self.dims = (1, 28, 28)
channels, width, height = self.dims
self.transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.1307, ), (0.3081, )),
])
# Define PyTorch model
self.model = nn.Sequential(
nn.Flatten(),
nn.Linear(channels * width * height, hidden_size),
nn.ReLU(),
nn.Dropout(0.2),
nn.Linear(hidden_size, hidden_size),
nn.ReLU(),
nn.Dropout(0.2),
nn.Linear(hidden_size, hidden_size),
nn.ReLU(),
nn.Dropout(0.2),
nn.Linear(hidden_size, self.num_classes),
)
def forward(self, x):
# predict or inference return
x = self.model(x)
return F.log_softmax(x, dim=1)
def training_step(self, batch, batch_idx):
x, y = batch
logits = self(x)
loss = F.nll_loss(logits, y)
return {'loss': loss}
def validation_step(self, batch, batch_idx):
x, y = batch
logits = self(x)
loss = F.nll_loss(logits, y)
preds = torch.argmax(logits, dim=1)
acc = accuracy(preds, y)
# Calling self.log will surface up scalars for you in TensorBoard
self.log('val_loss', loss, prog_bar=True)
self.log('val_acc', acc, prog_bar=True)
return {'val_loss': loss, 'val_acc': acc}
def test_step(self, batch, batch_idx):
# Here we just reuse the validation_step for testing
return self.validation_step(batch, batch_idx)
def configure_optimizers(self):
optimizer = torch.optim.Adam(self.parameters(), lr=self.learning_rate)
return optimizer
def prepare_data(self):
# download
MNIST(self.data_dir, train=True, download=True)
MNIST(self.data_dir, train=False, download=True)
def setup(self, stage=None):
if stage == 'fit' or stage is None:
mnist_full = MNIST(self.data_dir, train=True, download=False, transform=self.transform)
self.mnist_train, self.mnist_val = random_split(mnist_full, [55000, 5000])
if stage == 'test' or stage is None:
self.mnist_test = MNIST(self.data_dir, train=False, download=False, transform=self.transform)
def train_dataloader(self):
return DataLoader(self.mnist_train, batch_size=BATCH_SIZE)
def val_dataloader(self):
return DataLoader(self.mnist_val, batch_size=BATCH_SIZE)
def test_dataloader(self):
return DataLoader(self.mnist_test, batch_size=BATCH_SIZE)
조금 코드를 코쳐서 모델을 만들어 보았다.
GPU 2장을 사용할 건데, 이걸 PyTorch 에서는 Distributed mode 라고 명명한다.
문서에서 말하는 설명에 따르면,
prepare_data (things to do on 1 GPU/TPU not on every GPU/TPU in distributed mode).
setup (things to do on every accelerator in distributed mode).
즉,
데이터 다운로드가 필요한 상황에서 생각하면 좋겠다.
보통은 데이터를 미리 전처리를 거쳐서 학습에 feed 해주는 경우가 대부분이기 때문에
데이터 전처리를 마치고 AWS S3 에 존재하는 데이터라고 가정하고 생각하면 될 듯.
위에서 forward 는 기존(pure) PyTorch 에서 학습 영역으로 활용했는데,
PyTorch Lightning 에서는 forward를 prediction or inference 정도의 역할을 주고,
training_step 에서 loss를 반환해주면서 모델의 레이어들을 학습시킨다고 생각하면 된다.
이제 학습을 진행해 볼껀데,
진짜 GPU 2장에 모두 돌아가는지,
학습이 잘 진행되는지!
model = LitMNIST()
trainer = Trainer(
gpus=AVAIL_GPUS,
accelerator="dp",
max_epochs=10,
progress_bar_refresh_rate=10,
)
trainer.fit(model)
fit ? 파이썬 scikit-learn or keras 에서 볼 수 있는 함수다. 반갑다.
**혹시 왜 GPU 2장의 리소를 다 사용하지도 않으면서 굳이 2장을 모두 사용하는 이유는
그냥. 하고싶어서. 자랑하려고. 아니다.
모델, 데이터 크기가 커지면 Multi-GPU를 사용해서 batch_size 제약에서 조금씩 자유로워 질 수 있다.
accelerator="dp" 파라미터에서 dp 는 2개의 GPU 에서(Single-Node) 분산 학습을 진행하겠다는 선언이다.
ddp 라는 옵션도 있는데 이는 Multi-GPU & Multi-Node, 즉 다양한 분산 컴퓨터 시스템에서 다양한 GPU 를 함께 사용 할 때 필요하다.
Epoch 10 번까지 진행한다. GPU 라서 빠르다.
여기서 놀라운 사실은,
학습 로그가 자동으로 ./lightning_logs 라는 곳에 저장된다.
즉, 같은 위치에서 tensorboard --logdir ./lightning_logs/ --host 0.0.0.0 --port 8899 명령을 주면
바로 시각적인 학습 진행 상태를 확인 할 수 있다.
마지막 확인 사항은 모델 저장 / 불러오기
이건 정말... 쉽다.
모델 클래스를 잘 만들었다면(학습, 예측, 도출 등을 잘 정리해야 가능)
trainer.save_checkpoint("./mnist_pl_model.ckpt") # 저장
loaded_model = LitMNIST.load_from_checkpoint("./mnist_pl_model.ckpt") # 불러오기
쉽죠?
기존 model 이라는 변수가 있고, 새로운 loaded_model 이라는 변수가 있는데,
과연 두 모델은 같은 성능을 내는지 확인이 필요하다.
model.eval()
model.freeze()
loaded_model.eval()
loaded_model.freeze()
random_val = torch.rand((1, 28, 28)) # input_data 타입을 맞춰서 하나 만들고,
print(model(random_val))
print(loaded_model(random_val))
해보면 같은 결과값 (10개의 output_shape 에 맞춰서) 를 확인 할 수 있다.
이제 성능적인 면을 확인해보면,
위에서 만들어 두었던 trainer 를 사용해서 test를 진행한다.
trainer.test(model)
trainer.test(loaded_model)
모델을 넘겨주면 기존 모델 클래스에 존재하는 test_step 을 진행할 수 있다.
꼼꼼하게 확인해보았는데, 같은 숫자로 판명됐다.
복잡한 모델 구조를 편리하게 만들 수 있는 노하우(?)는 차차 생길 것이다.
과감한 결단이 필요했다.
Tensorflow 만 사용하다가 PyTorch 로 바꾸게된 이유는 생산성도 있지만, 함께 일하는 동료들과 딥러닝에 대한 커뮤니케이션을 높이기 위함이 더 컸다.
새로운 것들에 대한 호기심이 아직 내 안에 존재하고 있다는 것도 확인했다.
