[pytorch 튜토리얼] 2. 모델 구현 및 저장

1. 모델 매개변수 최적화하기

기존에 작성한 모델을 사용해 Fashion Data를 학습해볼 것임.

매개변수를 최적화하여 모델을 학습하고 테스트함.

학습 과정은 반복적이며 매 횟수를 epoch이라 부르고 추측과 정답의 loss값을 계산해 역전파하여 학습함(이 때 파라미터가 최적화optimize됨)

 

1.1 기본(pre-requisite) 코드

데이터셋 클래스 생성

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

training_data = datasets.FashionMNIST(
    root='data',
    train=True,
    download=True,
    transform=ToTensor()
)

test_data = datasets.FashionMNIST(
    root='data',
    train=False,
    download=True,
    transform=ToTensor()
)

train_dataloader = DataLoader(training_data, batch_size=64)
test_dataloader = DataLoader(test_data,batch_size=64)

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),
            nn.ReLU()
        )
    def forward(self,x):
        x = self.flatten(x)
        logits = self.linear_relu_stack(x)
        return logits

model = NeuralNetwork()

1.2 하이퍼파라미터(Hyperparameter)

하이퍼파라미터는 모델 최적화 과정을 제어할 수 있는 매개변수임.

learning_rate = 1e-3
batch_size = 64
epochs = 5

1.3 최적화 단계 (Optimization Loop)

이렇게 하이퍼파라미터를 설정한 뒤에는 최적화 단계를 통해 모델을 학습하고 최적화할 수 있음. 

하나의 에폭은 다음 두 부분으로 구성됨

• 학습 단계(train loop) - 학습용 데이터셋을 반복(iterate)하고 최적의 매개변수로 수렴함
• 검증/테스트 단계(validation/test loop) - 모델 성능이 개선되고 있는지를 확인하기 위해 테스트 데이터셋을 반복함(iterate)

학습 단계(training loop)에서 일어나는 몇 가지 개념들을 간략히 살펴보면 다음과 같음.

1.4 손실 함수(loss function)

학습할 때는 실제 값과 모델의 예측 값의 차이 정도를 측정해 이 값을 최소화 하는 것이 목표임.

손실 함수의 종류에는

1. 회귀 문제(regression task)에서는 nn.MSELoss(평균 제곱 오차)
2. 분류에서는 nn.NLLLoss(Negatice Log Likelihood)
3. nn.LogSoftmax와 nn.NLLLoss를 합친 nn.CrossEntropyLoss

등이 있음.

# 손실 함수를 초기화합니다.
loss_fn = nn.CrossEntropyLoss()

1.5 옵티마이저(Optimizer)

최적화 알고리즘은 모델 오류를 줄이기 위한 방식을 정의함 (예: SGD;Stochastic Gradient Descent(확률적 경사하강법))

모든 최적화 logic은 optimizer 객체에 encapsulate(캡슐화)됨. 이 예제에서는 SGD 옵티마이저를 사용하고 있지만, ADAM이나 RMSProp과 같은 옵티마이저는 다른 종류의 모델에서 적절하게 사용되니 옵티마이저에 대해서도 공부하도록.

optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate)

 

학습 단계(loop)에서 최적화는 세단계로 이루어짐.

• optimizer.zero_grad()를 호출하여 모델 매개변수의 변화도를 재설정함. 기본적으로 변화도는 더해지기 때문에 중복 계산을 막기 위해 반복할 때마다 명시적으로 0으로 설정함.
• loss.backward()를 호출해 prediction loss를 역전파함. PyTorch는 각 매개변수에 대한 손실의 변화도를 저장함.
• 변화도를 계산한 뒤에는 optimizer.step()을 호출하여 역전파 단계에서 수집된 변화도로 매개변수를 조정함.

1.6 전체 구현

def train_loop(dataloader, model, loss_fn, optimizer):
    size = len(dataloader.dataset)
    for batch, (X, y) in enumerate(dataloader):
        # 예측(prediction)과 손실(loss) 계산
        pred = model(X)
        loss = loss_fn(pred,y)

        # 역전파
        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 test_loop(dataloader, model, loss_fn):
    size = len(dataloader.dataset)
    num_batches = len(dataloader)
    test_loss, correct = 0, 0
    
    with torch.no_grad():
        for X, y in dataloader:
            pred = model(X)
            test_loss += loss_fn(pred, y).item()
            correct += (pred.argmx(1) ==y).type(torch.float).sum().item()
    test_loss /= num_batches
    correct /= size
    print(f"Test Error: \n Accuracy: {(100*correct):>0.1f}%, Avg loss: {test_loss:>8f} \n")

train_loop와 test_loop는 optimizing을 반복하는 코드임

 

2. 모델 저장하고 불러오기

import torch
import torch.onnx as onnx
import torchvision.models as models

 

2.1 모델 가중치 저장하고 불러오기

pytorch 모델은 학습한 매개변수를 state_dict라고 불리는 내부 상태 사전(internal state dictionary)에 저장함. 이 상태 값들은 torch.save 메소드를 사용해 저장할 수 있음.

model = models.vgg16(pretrained=True)
torch.save(model.state_dict(), 'model_weights.pth')

모델 가중치를 불러오기 위해서는 동일한 모델의 인스턴스(instance)를 생성한 다음에 load_state_dict() 메소드를 사용하여 매개변수들을 불러옴.

model = models.vgg16() # 기본 가중치를 불러오지 않으니 pretraine = True 설정 x
model.load_state_dict(torch.load('model_weights.pth'))
model.eval()

2.2 모델 형태를 포함해 저장하고 불러오기

torch.save(model, 'model.pth')

이처럼 모델 불러올 수 있음

model = torch.load('model.pth')

2.3 모델을 ONNX로 내보내기

input_image = torch.zeros((1, 3, 224, 224))
onnx.export(model, input_image, 'model.onnx')

Pytorch 실행 그래프의 동적 특성(dynamic nature)때문에 내보내는 과정에 ONNX 모델을 생성하기 위해 실행 그래프를 탐색(travers)해야 함. 이러한 이유 때문에 내보내기 단계에서는 적절한 크기의 테스트 변수를 전달해야함.