6

[Space-S x KaKR] 그래프 러닝 및 해커톤

SPACE-S

파이토치에서 과적합 전의 최적 모델을 저장하고 불러오는 방법

AF 김태영
2022.09.27 13:06
563

현재 제공된 베이스라인을 살펴보시면 전체 에포크만큼 학습한 최종 모델로 제출 파일을 생성합니다. 이러면 과적합이 된 모델로 추론했을 가능성이 높습니다. (과적합에 대해서 더 알고싶으시다면, 이 태스크로 >> 다층 퍼셉트론 신경망 검증하기) 과적합 전의 최적 모델의 결과를 얻기 위해 파이토치에서 학습한 모델을 저장하고 불러오는 방법에 대해서 설명드리고자 합니다. 이 방법은 여러가지 목적으로 사용되므로 알아두시면 요긴하게 사용할 수 있습니다.

  • 예를들어, 학습이 오래 걸릴 경우, 매 에포크마다 저장해놓았다가, 원할 때 다시 로딩해서 사용할 수 있습니다. 마치 긴 문서 작업이나 어드벤처 게임할 때 중간중간 저장하는 것 처럼요. 
  • 실제로 모델을 사용할때도 매번 학습해서 사용하는 것이 아니라 한 번 학습한 모델을 추론환경에서 사용하기 때문에 이때도 저장 및 불러오는 기능이 필요하겠죠?
  • 하지만 이번에는 과적합이 되기 전의 최적의 모델을 획득하기 위한 목적으로 사용해보겠습니다. 즉 매 에포크마다 검증 손실이 가장 낮은 것이 최적의 모델이라고 가정하고, 학습과정에서는 현재까지 에포크에서 가장 낮은 검증 손실이 확인될 때마다 모델 파일을 갱신합니다. 추론과정에서는 저장된 모델이 가장 낮은 검증손실을 가졌으므로, 이를 불러와서 사용합니다.

모델 저장 및 불러오는 코드 예시를 설명하기 위해 파이토치 공식 튜토리얼 중 퀵스타트 코드에 적용해보겠습니다.

필요한 패키지 선언

파이토치 모델 및 데이터셋을 사용하기 위해 필요한 패키지를 불러옵니다.

import torch
from torch import nn
from torch.utils.data import DataLoader
from torchvision import datasets
from torchvision.transforms import ToTensor

데이터셋 정의

Fashion MNIST 데이터셋을 사용합니다. 이 데이터셋에 대해서 더 알고 싶으시면, 아래 태스크를 살펴보시기 바랍니다.

* [다층 퍼셉트론 신경망 기반 패션아이템 이미지 분류](https://aifactory.space/class/detail/1674)

# 훈련셋 설정
training_data = datasets.FashionMNIST(
   root="data",
   train=True,
   download=True,
   transform=ToTensor(),
)
# 검증셋 설정
val_data = datasets.FashionMNIST(
   root="data",
   train=False,
   download=True,
   transform=ToTensor(),
)
# 데이터로더 설정
train_dataloader = DataLoader(training_data, batch_size=64)
val_dataloader = DataLoader(val_data, batch_size=64)

모델 정의

nn.Module을 상속받아 간단한 모델을 정의합니다. 이어서 손실함수 및 최적화기도 함께 정의합니다.

# 아키텍처 정의
class NeuralNetwork(nn.Module):
   def __init__(self):
       super(NeuralNetwork, self).__init__()
       self.flatten = nn.Flatten()
       self.linear_relu_stack = nn.Sequential(
           nn.Linear(28*28, 512),
           nn.ReLU(),
           nn.Linear(512, 512),
           nn.ReLU(),
           nn.Linear(512, 10)
       )
   def forward(self, x):
       x = self.flatten(x)
       logits = self.linear_relu_stack(x)
       return logits
# 모델 생성
device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Using {device} device")
model = NeuralNetwork().to(device)
loss_fn = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr = 0.01)

학습 및 검증 함수 정의

머신러닝에서는 크게 학습과정과 추론과정으로 나누어집니다. 학습과정은 데이터셋으로 모델의 가중치를 업데이트하는 것이고, 추론과정은 학습 모델을 사용하는 것입니다. 모델 아키텍처 특성에 따라 학습과 추론 시 다르게 동작하는 경우도 있기 때문에 학습 과정에서는 'model.train()'을 추론과정에서는 'model.eval()'을 명시적으로 호출합니다.

학습과정에서는 모델의 가중치를 업데이트하기 때문에 손실을 계산하는 부분과 계산된 손실로 역전파하는 코드가 포함되어 있습니다. 반대로 추론 과정에서는 역전파를 하지 않아 그래이디언트를 계산할 필요가 없으므로 torch.no_grad() 함수를 호출합니다.

모델 검증 및 과적합을 확인하기 위해서는 매 에포크마다 검증셋으로 모델 평가를 수행합니다. 모델 평가시에는 이미 정답을 알고 있는 데이터로 모델을 추론하여 모델 결과와 정답으로 손실을 계산합니다. 이때 모델 가중치는 업데이트 하지 않습니다.

# 학습 함수
def train(dataloader, model, loss_fn, optimizer):
   size = len(dataloader.dataset)
   model.train()
   for batch, (X, y) in enumerate(dataloader):
       X, y = X.to(device), y.to(device)
       # Compute prediction error
       pred = model(X)
       loss = loss_fn(pred, y)
       # Backpropagation
       optimizer.zero_grad()
       loss.backward()
       optimizer.step()
       if batch % 100 == 0:
           loss, current = loss.item(), batch * len(X)
           print(f"loss: {loss:>7f}  [{current:>5d}/{size:>5d}]")
# 검증 함수
def val(dataloader, model, loss_fn):
   size = len(dataloader.dataset)
   num_batches = len(dataloader)
   model.eval()
   val_loss, correct = 0, 0
   with torch.no_grad():
       for X, y in dataloader:
           X, y = X.to(device), y.to(device)
           pred = model(X)
           val_loss += loss_fn(pred, y).item()
           correct += (pred.argmax(1) == y).type(torch.float).sum().item()
   val_loss /= num_batches
   correct /= size
   print(f"Val Error: \n Accuracy: {(100*correct):>0.1f}%, Avg loss: {val_loss:>8f} \n")
   return val_loss        

학습 및 검증 수행

학습 시에 매 에포크마다 검증셋으로 검증손실을 구합니다. 검증손실이 가장 낮을 때마다 "best_model.pth" 이란 이름으로 파이토치 모델을 저장합니다. 같은 파일명으로 덮어쓰기 때문에 이 파일명으로 저장된 모델은 지정된 에포크에서 가장 낮은 검증 손실을 가지고 있습니다.

epochs = 30
best_val_loss = 999.0
for t in range(epochs):
   print(f"Epoch {t+1}\n-------------------------------")
   train(train_dataloader, model, loss_fn, optimizer)
   curr_val_loss = val(val_dataloader, model, loss_fn)
   if curr_val_loss < best_val_loss:
     print(f"The best model is saved. epoch: {t+1} val_loss: {curr_val_loss}\n")
     torch.save(model.state_dict(), "best_model.pth")
     best_val_loss = curr_val_loss
print("Done!")

모델 불러오기

model을 선언한 뒤, "best_model.pth" 파일에서 가중치를 가지고 온 뒤 model에 설정합니다.

# 모델 선언 및 모델 불러오기
model = NeuralNetwork()
model.load_state_dict(torch.load("best_model.pth"))

추론하기

최적의 모델로 임의의 데이터에 대한 추론을 수행합니다.

# 불러온 모델으로 추론한다.

classes = [
   "T-shirt/top",
   "Trouser",
   "Pullover",
   "Dress",
   "Coat",
   "Sandal",
   "Shirt",
   "Sneaker",
   "Bag",
   "Ankle boot",
]
model.eval()
x, y = val_data[0][0], val_data[0][1]
with torch.no_grad():
   pred = model(x)
   predicted, actual = classes[pred[0].argmax(0)], classes[y]
   print(f'Predicted: "{predicted}", Actual: "{actual}"')

참고

  • https://pytorch.org/tutorials/beginner/basics/quickstart_tutorial.html
6
3개의 댓글
로그인 후 이용해주세요!