딥 러닝 모델을 개발할 때에는 모델의 학습이 완료된 후, 이를 저장하고 필요할 때 로드하여 사용해야 할 때가 있습니다. 이러한 저장 및 로딩은 모델을 재학습하는데 많은 시간을 절약할 수 있습니다. 이번 포스트에서는 PyTorch에서 모델을 저장하고 로드하는 방법에 대해 알아보겠습니다.

모델 저장하기

모델을 저장하기 위해서는 먼저 모델의 가중치와 매개변수를 저장할 파일 경로를 정해야 합니다. 이후 torch.save() 함수를 사용하여 파라미터를 저장합니다.

import torch
import torch.nn as nn

# 예시 모델
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.fc = nn.Linear(10, 5)

    def forward(self, x):
        x = self.fc(x)
        return x

model = Net()

# 모델의 가중치 저장하기
torch.save(model.state_dict(), 'model_weights.pth')

위 코드에서는 예시로 간단한 10차원 입력을 받는 선형 레이어를 갖는 신경망 모델을 생성하였습니다. 모델의 가중치를 저장하기 위해 model.state_dict() 함수를 사용했습니다. 이 함수는 모델의 가중치를 Python 사전(Dictionary)으로 반환합니다. 이후 반환된 가중치 사전을 torch.save() 함수를 사용하여 파일로 저장합니다.

모델 로드하기

모델을 로드하기 위해서는 먼저 모델을 정의해야 합니다. 그리고 torch.load() 함수를 사용하여 가중치를 읽어옵니다.

import torch
import torch.nn as nn

# 예시 모델
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.fc = nn.Linear(10, 5)

    def forward(self, x):
        x = self.fc(x)
        return x

model = Net()

# 모델의 가중치 로드하기
model.load_state_dict(torch.load('model_weights.pth'))

위 코드에서는 이전에 저장한 가중치를 torch.load() 함수를 사용하여 읽어온 후, model.load_state_dict() 함수를 통해 모델에 로드합니다. 이제 모델은 이전에 저장된 가중치를 가지게 됩니다.

전체 모델 저장하기

가중치 뿐만 아니라 전체 모델을 저장하고 싶을 때도 있습니다. 이 경우에는 모델 인스턴스를 그대로 저장하면 됩니다.

import torch
import torch.nn as nn

# 예시 모델
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.fc = nn.Linear(10, 5)

    def forward(self, x):
        x = self.fc(x)
        return x

model = Net()

# 모델 전체 저장하기
torch.save(model, 'model.pth')

위 코드에서는 torch.save() 함수를 사용하여 전체 모델을 model.pth 파일로 저장합니다. 이 때에는 모델의 가중치, 구조, 그리고 설정 등이 모두 저장됩니다. 모델을 로딩할 때에도 torch.load() 함수를 사용하면 됩니다.

결론

PyTorch에서는 모델의 가중치, 전체 모델을 저장하고 로드하는 과정이 간단합니다. 저장과 로딩을 통해 모델을 손쉽게 재사용할 수 있습니다. 이를 통해 중요한 모델을 보다 효율적으로 관리하고 활용할 수 있습니다.