pytorch는 가장 널리 쓰이는 딥러닝 프레임워크중 하나이죠. 딥러닝, 머신러닝을 공부하시는 분이라면 들어보셨을겁니다.
써보신 분들은 아시겠지만 pytorch 자체가 굉장히 자유도가 높은 프레임워크이죠. 일반적으로 딥러닝 모델을 구현할 경우 크게 복잡하지는 않지만 GPU, 분산 처리 등을 구현할 경우 코드가 조금씩 복잡해집니다. pytorch lightning은 굉장히 잘 정돈된 하나의 틀을 제공해줍으로써 pytorch의 가독성을 높여줍니다.
이제 본격적으로 pytorch lightning에 대해서 알아보도록 하죠.
pytorch lightning 사용 이유
1. 간단한 학습, 검증과정
pytorch lightning은 pytorch의 학습, 검증 과정을 training_step, validation_step 등의 함수로 제공합니다. 이미 틀이 있기 때문에 내부만 채워주시면 아주 간결하게 실행이 가능합니다.
2. 통일된 구조
위와 동일한 맥락입니다. 구조가 통일되어있기 때문에 가독성이 높습니다.
3. 자동화된 기능
logging, checkpoint, 병렬처리 등이 자동화 되어있습니다. 이 또한 가독성을 높여줍니다.
이렇게 말씀드리면 느낌이 잘 안오실것 같습니다.
간단하게 코드로 비교해봅시다.
pytorch에 비해 굉장히 간결해졌죠?
pytorch lightning 사용 방법
이번 장에서는 pytorch lightning을 어떻게 사용하는지 mnist예제로 한번 살펴보도록 하겠습니다.
1. 설치
pip install pytorch-lightning
2. 모델 정의
먼저 MNIST데이터셋을 분류하는 pytorch 코드를 먼저 보시죠. 이 코드를 pytorch lightning으로 변환시켜보겠습니다.
import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms
from torch.utils.data import DataLoader
# 데이터 변환
transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,), (0.5,))])
# 데이터 로딩
train_dataset = datasets.MNIST(root='./data', train=True, transform=transform, download=True)
test_dataset = datasets.MNIST(root='./data', train=False, transform=transform)
train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=64)
# 모델 정의
class Net(nn.Module):
def __init__(self):
super(Net, self).__init__()
self.fc1 = nn.Linear(28*28, 128)
self.fc2 = nn.Linear(128, 10)
def forward(self, x):
x = x.view(-1, 28*28)
x = torch.relu(self.fc1(x))
x = self.fc2(x)
return x
# 모델, 손실 함수, 옵티마이저 설정
model = Net()
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.01)
# 훈련 루프
for epoch in range(10):
model.train()
for data, target in train_loader:
optimizer.zero_grad()
output = model(data)
loss = criterion(output, target)
loss.backward()
optimizer.step()
print(f"Epoch {epoch + 1}: Loss {loss.item()}")
# 모델 평가
model.eval()
correct = 0
total = 0
with torch.no_grad():
for data, target in test_loader:
output = model(data)
_, predicted = torch.max(output, 1)
total += target.size(0)
correct += (predicted == target).sum().item()
print(f"Accuracy: {100 * correct / total}%")
pytorch 코드를 보시면 데이터 정의, 모델 정의, 학습 정의를 모두 따로 정의해 줍니다.
pytorch lightning 같은 경우에는 클래스 내에 한번에 정의해줍니다.
클래스는 pytorch_lightning.LightningModule를 상속받습니다.
기본 틀은 아래와 같습니다.
class MNISTClassifier(pl.LightningModule):
def __init__(self):
super().__init__()
#### layer들 정의해주기
def forward(self, x):
#### network 정의해주기
def configure_optimizers(self):
#### optimizer 정의해주기
def train_dataloader(self):
#### train data 불러와주기
def val_dataloader(self):
#### val data 불러와주기
def test_dataloader(self):
#### val data 불러와주기
def training_step(self, batch, batch_idx):
#### loss 계산해주기(정의해주기)
def validation_step(self, batch, batch_idx):
#### 매 batch 마다 실행되는 함수
#### batch의 loss계산, acc계산등을 수행
def on_validation_epoch_end(self, outputs):
#### 매 epoch마다 실행되는 함수
#### validation_step에서 모은 loss, acc등을 모아서 평균내기 등
위의 코드 주석으로도 설명을 간단하게 달았지만, 여기서도 한번씩 짚고 가죠!
init과 forward는 기존과 동일하기 때문에 설명을 생략하겠습니다.
configure_optimizers
optimizer를 정의해 주는 부분입니다. 단순히 정의해준 뒤에 return시켜주면됩니다.
def configure_optimizers(self):
return optim.SGD(self.parameters(), lr=0.01)
train_dataloader, val_dataloader, test_dataloader
로드한 데이터를 DataLoader에 태워줍니다.
def train_dataloader(self):
transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,), (0.5,))])
train_dataset = datasets.MNIST(root='./data', train=True, transform=transform, download=True)
return DataLoader(train_dataset, batch_size=64, shuffle=True)
def validation_dataloader(self):
transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,), (0.5,))])
test_dataset = datasets.MNIST(root='./data', train=False, transform=transform)
return DataLoader(test_dataset, batch_size=64)
training_step, validation_step, test_step
batch 기준으로 작동하는 함수들입니다. 각각 학습, 검증, 테스트에 대한 loss, 성능을 계산해줍니다.
def training_step(self, batch, batch_idx):
data, target = batch
output = self.forward(data)
loss = nn.functional.cross_entropy(output, target)
self.log('train_loss', loss)
return loss
def validation_step(self, batch, batch_idx):
data, target = batch
output = self.forward(data)
loss = nn.functional.cross_entropy(output, target)
pred = output.argmax(dim=1, keepdim=True)
correct = pred.eq(target.view_as(pred)).sum().item()
self.log('val_loss', loss)
self.log('val_acc', correct / len(data))
return {'val_loss':loss, 'val_acc':correct / len(data)}
on_validation_epoch_end
한 epoch의 모든 batch validation_step의 결과값을 입력받습니다.
보통 전체 성능 계산시 사용합니다.
위에서 결과값으로 정의한 딕셔너리를 그대로 list로 쌓아서 받는다 생각하시면 됩니다.
def on_validation_epoch_end(self, outputs):
avg_loss = torch.stack([x['val_loss'] for x in outputs]).mean()
avg_acc = torch.stack([x['val_acc'] for x in outputs]).mean()
self.log('avg_val_loss', avg_loss)
self.log('avg_val_acc', avg_acc)
이제 모든 정의가 끝났다면 fit해주시면 됩니다.
3. 학습
## 모델 정의
model = MNISTClassifier()
## trainer 정의
trainer = pl.Trainer(max_epochs=10)
## trainer 태우기
trainer.fit(model)
학습은 정말 간단하죠?
오늘은 대략적인 pytorch lightning 설명과 사용법을 다루어보았는데요.
심화 내용은 따로 포스팅을 하겠습니다.
4. 전체 코드
import pytorch_lightning as pl
import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms
from torch.utils.data import DataLoader
class MNISTClassifier(pl.LightningModule):
def __init__(self):
super().__init__()
self.fc1 = nn.Linear(28*28, 128)
self.fc2 = nn.Linear(128, 10)
def forward(self, x):
x = x.view(-1, 28*28)
x = torch.relu(self.fc1(x))
x = self.fc2(x)
return x
def configure_optimizers(self):
return optim.SGD(self.parameters(), lr=0.01)
def training_step(self, batch, batch_idx):
data, target = batch
output = self.forward(data)
loss = nn.functional.cross_entropy(output, target)
self.log('train_loss', loss)
return loss
def train_dataloader(self):
transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,), (0.5,))])
train_dataset = datasets.MNIST(root='./data', train=True, transform=transform, download=True)
return DataLoader(train_dataset, batch_size=64, shuffle=True)
def validation_dataloader(self):
transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,), (0.5,))])
test_dataset = datasets.MNIST(root='./data', train=False, transform=transform)
return DataLoader(test_dataset, batch_size=64)
def validation_step(self, batch, batch_idx):
data, target = batch
output = self.forward(data)
loss = nn.functional.cross_entropy(output, target)
pred = output.argmax(dim=1, keepdim=True)
correct = pred.eq(target.view_as(pred)).sum().item()
self.log('val_loss', loss)
self.log('val_acc', correct / len(data))
return {'val_loss':loss,'val_acc':correct/ len(data)}
def on_validation_epoch_end(self, outputs):
avg_loss = torch.stack([x['val_loss'] for x in outputs]).mean()
avg_acc = torch.stack([x['val_acc'] for x in outputs]).mean()
self.log('avg_val_loss', avg_loss)
self.log('avg_val_acc', avg_acc)
model = MNISTClassifier()
trainer = pl.Trainer(max_epochs=10)
trainer.fit(model)
'인공지능공부 > 프레임워크' 카테고리의 다른 글
pytorch lightning multi-gpu사용하기 (0) | 2023.09.12 |
---|