[Pytroch] Check point 저장하고 불러오기

추론(inference) 또는 학습(training)의 재개를 위해 체크포인트(checkpoint) 모델을 저장하고 불러오는 것은 마지막으로 중단했던 부분을 선택하는데 도움을 줄 수 있다. 체크포인트를 저장할 때는 단순히 모델의 state_dict 이상의 것을 저장해야 한다. 모델 학습 중에 갱신되는 퍼버와 매개변수들을 포함하는 옵티마이저(Optimizer)의 state_dict를 함께 저장하는 것이 중요하다. 이 외에도 중단 시점의 에포크(epoch), 마지막으로 기록된 학습 오차(training loss), 외부 torch.nn.Embedding 계층 등, 알고리즘에 따라 저장하고 싶은 항목들이 있을 것이다.

개요

여러 체크포인트들을 저장하기 위해서는 사전(dictionary)에 체크포인트들을 구성하고 torch.save() 를 사용하여 사전을 직렬화(serialize)해야 한다. 일반적인 PyTorch에서는 이러한 여러 체크포인트들을 저장할 때 .tar 확장자를 사용하는 것이 일반적인 규칙이다. 항목들을 불러올 때에는, 먼저 모델과 옵티마이저를 초기화하고, torch.load()를 사용하여 사전을 불러온다. 이후 원하는대로 저장한 항목들을 사전에 조회하여 접근할 수 있다.

 

여러 체크포인트들을 어떻게 저장하고 불러오는지 살펴보자.

단계(Steps)

1. 데이터 불러올 때 필요한 라이브러리들 불러오기
2. 신경망을 구성하고 초기화하기
3. 옵티마이저 초기화하기
4. 일반적인 체크포인트 저장하기
5. 일반적인 체크포인트 불러오기

1. 데이터 불러올 때 필요한 라이브러리들 불러오기

import torch
import torch.nn as nn
import torch.optim as optim

2. 신경망을 구성하고 초기화하기

이미지를 학습하는 신경망을 만들어자.

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(3, 6, 5)
        self.pool = nn.MaxPool2d(2, 2)
        self.conv2 = nn.Conv2d(6, 16, 5)
        self.fc1 = nn.Linear(16 * 5 * 5, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)

    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        x = x.view(-1, 16 * 5 * 5)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x

net = Net()
print(net)

3. 옵티마이저 초기화하기

모멘텀(momentum)을 갖는 SGD를 사용하자.

optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)

4. 일반적인 체크포인트 저장하기

관련된 모든 정보들을 모아서 사전을 구성한다.

# 추가 정보
EPOCH = 5
PATH = "model.pt"
LOSS = 0.4

torch.save({
            'epoch': EPOCH,
            'model_state_dict': net.state_dict(),
            'optimizer_state_dict': optimizer.state_dict(),
            'loss': LOSS,
            }, PATH)

 

5. 일반적인 체크포인트 불러오기

모델과 옵티마이저를 먼저 초기화한 뒤, 사전을 불러오는 것을 기억하자.

model = Net()
optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)

checkpoint = torch.load(PATH)
model.load_state_dict(checkpoint['model_state_dict'])
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
epoch = checkpoint['epoch']
loss = checkpoint['loss']

model.eval()
# - 또는 -
model.train()

 

참고: https://tutorials.pytorch.kr/recipes/recipes/saving_and_loading_a_general_checkpoint.html