[PyTorch] 모델 저장하기 및 불러오기

PyTorch를 사용하다보면, 학습 결과를 저장한 뒤, 불러와야 하는 경우가 생긴다.
이러한 상황을 위해 PyTorch로 만든 모델을 저장하고, 불러오는 방법에 대해 알아보자.

---

state_dict란?

각 layer의 파라미터에 대한 정보가 들어있는 dict 타입의 객체이다.

# 모델 정의
class TheModelClass(nn.Module):
    def __init__(self):
        super(TheModelClass, 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

# 모델 초기화
model = TheModelClass()

# 모델의 state_dict 출력
print("Model's state_dict:")
for param_tensor in model.state_dict():
    print(param_tensor, "\t", model.state_dict()[param_tensor].size())

# 출력: 
# Model's state_dict:
# conv1.weight     torch.Size([6, 3, 5, 5])
# conv1.bias   torch.Size([6])
# conv2.weight     torch.Size([16, 6, 5, 5])
# conv2.bias   torch.Size([16])
# fc1.weight   torch.Size([120, 400])
# fc1.bias     torch.Size([120])
# fc2.weight   torch.Size([84, 120])
# fc2.bias     torch.Size([84])
# fc3.weight   torch.Size([10, 84])
# fc3.bias     torch.Size([10])

 

torch.save

• 객체를 Python의 pickle을 사용해 저장한다.
• 모델, 모델의 파라미터, state_dict를 저장할 수 있다.

# 모델의 파라미터를 저장
torch.save(model.state_dict(), PATH)

# 모델의 architecture와 함께 저장
torch.save(model, PATH)

 

torch.load

pickle을 사용해 저장되어 있는 객체 파일을 메모리에 올려 객체로 불러온다.

 

load_state_dict

모델의 파라미터를 불러온다.

# 모델의 파라미터만 load
model = TheModelClass()
model.load_state_dict(torch.load(state_dict가 저장된 경로))

# 모델의 architecture와 함께 load
model = torch.load(모델이 저장된 경로)

 

일반적으로 모델의 파라미터만 저장하고, 불러오는 방법을 권장한다.