[PyTorch] Checkpoint Model 저장 & 불러오기

⚠️ 본 내용은 PyTorch Korea의 공식 문서에 기반하여 공부한 내용을 적은것이니 양해바랍니다!

PyTorch에서 일반적인 체크포인트(checkpoint) 저장하기 & 불러오기

---

PyTorch에서 Inference & Training을 다시 하기 위해서 Checkpoint Model을 저장 & 불러오는것을 한번 해보겠습니다.

Intro

• PyTorch에서 여러 Checkpoint들을 저장하기 위해선 사전(Dictionary)에 Checkpoint들을 구성한후,
• torch.save()를 사용하여 Dictionary를 직렬화(Seralize)해야 합니다.
• PyTorch에서는 여러 체크포인트들을 저장할 때 .tar 확장자를 사용합니다.
• 항목들을 불러올 때에는, 먼저 모델과 옵티마이저를 초기화하고, torch.load()를 사용하여 사전을 불러옵니다.
• 이후 원하는대로 저장한 항목들을 사전에 조회하여 접근할 수 있습니다.

---

Setting (설정)

일단, 먼저 PyTorch 라이브러리를 설치해줘야 합니다.

!pip install torch

---

Model 저장 & 불러오는 단계

한번, 모델을 저장 & 불러오는 단계를 한번 구현해 보겠습니다.

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

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

---

Neural Network (신경망) 구성 & 초기화

예시를 들어 모델을 저장 & 불러오기 위한 신경망을 한번 구성해 보겠습니다.

class NeuarlNetwork(nn.Module):
    def __init__(self):
        super(NeuarlNetwork, self).__init__()
        self.conv1 = nn.Conv2d(3, 32, kernel_size=3, stride=1, padding=1)
        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1)
        self.pool = nn.MaxPool2d(kernel_size=2, stride=2, padding=0)
        self.fc1 = nn.Linear(64 * 8 * 8, 512)
        self.fc2 = nn.Linear(512, 64)
        self.fc3 = nn.Linear(64, 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, 64 * 8 * 8)
        x = F.relu(self.fc1(x))
        x = self.fc2(x)
        return x

net = NeuarlNetwork()
print(net)

NeuarlNetwork(
  (conv1): Conv2d(3, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
  (conv2): Conv2d(32, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
  (pool): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
  (fc1): Linear(in_features=4096, out_features=512, bias=True)
  (fc2): Linear(in_features=512, out_features=64, bias=True)
  (fc3): Linear(in_features=64, out_features=10, bias=True)
)

---

Optimizer 초기화 하기

Momentum(모멘텀)이랑 SGD를 사용하겠습니다.

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

---

Model의 Checkpoint 저장하기

Model에 관련된 정보를 불러와서 Dictionary를 구성해 보겠습니다.

# 추가 정보
EPOCH = 5  # 저장할 에폭 수
PATH = "model.pt"  # 모델 저장 경로
LOSS = 0.4  # 마지막 손실 값

# 모델과 옵티마이저 상태 저장
torch.save({
            'epoch': EPOCH,  # 현재 에폭 수
            'model_state_dict': net.state_dict(),  # 모델의 상태(dictionary)
            'optimizer_state_dict': optimizer.state_dict(),  # 옵티마이저의 상태(dictionary)
            'loss': LOSS,  # 마지막 손실 값
            }, PATH)  # 지정된 경로에 저장

---

Model의 Checkpoint 불러오기

Model과 Optimizer를 초기화 한후, Checkpoint를 저장한 Dictionary를 불러와야 합니다.

# 모델과 옵티마이저 초기화
model = Net()  # 사용자 정의 모델 인스턴스 생성
optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)  # SGD 옵티마이저 초기화

# 체크포인트 불러오기
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()  # 학습 모드로 설정 (드롭아웃, 배치 정규화 활성화)

• 추론(inference)을 실행하기 전에 model.eval() 을 호출하여 드롭아웃(dropout)과 배치 정규화 층(batch normalization layer)을 평가(evaluation) 모드로 바꿔야한다는 것을 기억해야합니다.
• 만약 이것을 빼먹으면 일관성 없는 추론 결과를 얻게 됩니다.
• 또한 학습을 계속하길 원한다면 model.train() 을 호출하여 이 층(layer)들이 학습 모드인지 확인(ensure)해야 합니다.