9-2. PyTorch 고급: 모델 저장 및 로드

1. 모델 저장과 로드의 중요성

딥러닝 모델은 학습 과정에서 수많은 파라미터를 조정하며 복잡한 구조를 형성합니다. 이러한 모델을 파일로 저장하고 필요할 때 다시 불러오는 것은 딥러닝 프로젝트의 필수적인 요소입니다. 모델 저장과 로드를 통해 우리는 다음과 같은 이점을 얻을 수 있습니다.

• 재사용성: 한 번 학습된 모델을 다른 작업이나 데이터에 재사용하여 학습 시간을 절약하고, 새로운 문제에 빠르게 적용할 수 있습니다.
• 배포: 학습된 모델을 실제 서비스 환경에 배포하여 예측 서비스를 제공할 수 있습니다.
• 실험 관리: 다양한 모델과 하이퍼파라미터 조합에 대한 실험 결과를 저장하고 관리하여, 가장 좋은 성능을 보이는 모델을 선택할 수 있습니다.
• 중단된 학습 재개: 학습 중단 시점부터 다시 학습을 시작하여 학습 시간 낭비를 줄이고, 모델 훈련의 유연성을 확보할 수 있습니다.

2. PyTorch에서의 모델 저장 및 로드 방법

PyTorch는 모델 저장 및 로드를 위한 다양한 기능을 제공합니다. 크게 모델의 전체 구조와 가중치를 함께 저장하는 방법, 그리고 가중치만 저장하는 방법으로 나눌 수 있습니다. 각 방법은 사용 목적과 상황에 따라 적절하게 선택되어야 합니다.

1) 모델 전체 저장 및 로드

모델의 구조, 가중치, optimizer 상태 등 모델과 관련된 모든 정보를 하나의 파일로 저장하는 방법입니다. torch.save() 함수를 사용하여 모델을 저장하고, torch.load() 함수를 사용하여 모델을 로드합니다.

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

# 예시 모델 정의
class SimpleNN(nn.Module):
    def __init__(self, input_size, hidden_size, output_size):
        super(SimpleNN, self).__init__()
        self.fc1 = nn.Linear(input_size, hidden_size)
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(hidden_size, output_size)

    def forward(self, x):
        out = self.fc1(x)
        out = self.relu(out)
        out = self.fc2(out)
        return out

# 모델, optimizer, loss function 정의
input_size = 10
hidden_size = 20
output_size = 1
model = SimpleNN(input_size, hidden_size, output_size)
optimizer = optim.Adam(model.parameters(), lr=0.001)
criterion = nn.MSELoss()

# 가짜 데이터 생성
inputs = torch.randn(16, input_size)
targets = torch.randn(16, output_size)

# 학습
for epoch in range(10):
    optimizer.zero_grad()
    outputs = model(inputs)
    loss = criterion(outputs, targets)
    loss.backward()
    optimizer.step()
    print(f'Epoch [{epoch+1}/10], Loss: {loss.item():.4f}')

# 모델 저장
model_path = 'simple_nn.pth'
torch.save({
    'epoch': 10,  # 마지막 epoch
    'model_state_dict': model.state_dict(),
    'optimizer_state_dict': optimizer.state_dict(),
    'loss': loss,
}, model_path)

위 코드에서 torch.save()는 딕셔너리 형태로 모델, optimizer, epoch, loss 등 필요한 모든 정보를 저장합니다. model_state_dict()는 모델의 가중치 정보를 담고 있으며, optimizer_state_dict()는 optimizer의 상태를 저장합니다. 이러한 정보들은 모델을 로드한 후 학습을 이어서 진행하거나, 모델의 상태를 복원하는 데 사용됩니다.

# 모델 로드
loaded_model_path = 'simple_nn.pth'
checkpoint = torch.load(loaded_model_path)
model = SimpleNN(input_size, hidden_size, output_size)
optimizer = optim.Adam(model.parameters(), lr=0.001)

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

model.eval() # 평가 모드로 변경
print(f'Loaded model from epoch: {epoch}, Loss: {loss.item():.4f}')

torch.load() 함수는 저장된 딕셔너리를 불러오고, model.load_state_dict()와 optimizer.load_state_dict() 함수를 사용하여 모델과 optimizer의 상태를 복원합니다. .eval() 함수는 모델을 평가 모드로 설정하여 드롭아웃(dropout)과 배치 정규화(batch normalization) 등의 동작을 비활성화합니다.

위의 이미지와 같이, 모델 전체 저장 방법은 모델 구조, 가중치, optimizer 상태를 하나의 파일에 묶어 저장합니다. 이 방법은 모델을 쉽게 배포하고, 학습을 중단한 시점부터 다시 시작하는 경우에 유용합니다. 하지만, 모델 구조가 변경되면 이전의 모델을 로드할 수 없다는 단점이 있습니다.

2) 가중치 (state_dict) 저장 및 로드

모델의 가중치 정보만 저장하고 로드하는 방법입니다. state_dict()는 모델의 각 레이어의 가중치와 bias를 포함하는 딕셔너리를 반환하며, load_state_dict() 함수를 사용하여 가중치를 로드합니다.

import torch
import torch.nn as nn

# 예시 모델 정의
class SimpleNN(nn.Module):
    def __init__(self, input_size, hidden_size, output_size):
        super(SimpleNN, self).__init__()
        self.fc1 = nn.Linear(input_size, hidden_size)
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(hidden_size, output_size)

    def forward(self, x):
        out = self.fc1(x)
        out = self.relu(out)
        out = self.fc2(out)
        return out

# 모델 정의
input_size = 10
hidden_size = 20
output_size = 1
model = SimpleNN(input_size, hidden_size, output_size)

# 가짜 데이터 생성 및 모델 학습 (생략)
# ...

# 가중치 저장
model_path = 'simple_nn_weights.pth'
torch.save(model.state_dict(), model_path)

torch.save() 함수에 model.state_dict()를 전달하여 가중치 딕셔너리를 저장합니다.

# 가중치 로드
loaded_model_path = 'simple_nn_weights.pth'
model = SimpleNN(input_size, hidden_size, output_size)
model.load_state_dict(torch.load(loaded_model_path))
model.eval()  # 평가 모드로 변경

# 모델 사용 (예: 예측)
# ...

모델을 로드하기 전에 동일한 모델 구조를 정의해야 합니다. 그 후 model.load_state_dict()를 사용하여 저장된 가중치를 로드합니다.

위의 그림은 가중치 저장 및 로드 과정을 보여줍니다. 이 방법은 모델 구조가 동일하다면, 다른 환경이나 다른 코드에서 학습된 모델의 가중치를 사용할 수 있다는 장점이 있습니다. 하지만 모델 구조가 변경되면 호환되지 않는다는 점에 유의해야 합니다.

3. 실용적인 고려 사항

1) 파일 경로 관리

모델 저장 및 로드 시, 파일 경로는 매우 중요합니다. 파일 경로는 절대 경로 또는 상대 경로로 지정할 수 있으며, 프로젝트 구조에 맞게 관리하는 것이 좋습니다.

• 상대 경로: 현재 작업 디렉토리를 기준으로 파일의 위치를 지정합니다. ('./models/my_model.pth')
• 절대 경로: 파일의 전체 경로를 지정합니다. ('/home/user/projects/my_model.pth')

프로젝트를 다른 환경으로 옮길 경우, 상대 경로를 사용하는 것이 더 유연합니다.

2) 장치 (CPU/GPU) 호환성

모델을 저장한 환경과 로드하는 환경의 장치(CPU 또는 GPU)가 다를 수 있습니다. 이 경우, torch.load() 함수의 map_location 인자를 사용하여 장치를 지정해야 합니다.

# GPU에서 저장된 모델을 CPU에서 로드
device = torch.device('cpu')
model = SimpleNN(input_size, hidden_size, output_size)
model.load_state_dict(torch.load(model_path, map_location=device))

map_location에 torch.device('cpu')를 지정하면, 저장된 텐서들이 CPU 메모리에 로드됩니다. GPU에서 모델을 로드하고 싶다면, torch.device('cuda') 또는 torch.device('cuda:0')과 같이 원하는 GPU 장치를 지정합니다. 만약 여러 개의 GPU를 사용하고 있다면, torch.device('cuda:0')과 같이 특정 GPU를 지정할 수 있습니다.

3) 버전 관리

모델을 저장할 때, 모델의 버전 정보를 함께 저장하는 것이 좋습니다. 이를 통해, 모델의 변경 사항을 추적하고, 필요한 경우 이전 버전의 모델을 로드할 수 있습니다.

torch.save({
    'model_state_dict': model.state_dict(),
    'version': '1.0',
    'date': '2023-10-27',
}, model_path)

version과 date와 같은 정보를 함께 저장하여 모델의 버전을 관리합니다.

4) 모델 아키텍처 변경

모델 구조가 변경된 경우, state_dict를 직접 수정하여 가중치를 로드하는 방법을 고려할 수 있습니다. 하지만, 이 방법은 모델 구조의 변경 정도에 따라 어려울 수 있으며, 주의해서 사용해야 합니다. 모델 구조가 크게 변경된 경우에는, 가중치를 처음부터 다시 학습하는 것이 더 안전할 수 있습니다.

4. 모델 저장 및 로드 시 발생 가능한 문제점과 해결 방법

1) FileNotFoundError

파일 경로가 잘못되었거나, 해당 경로에 파일이 없는 경우 발생합니다.

• 해결 방법: 파일 경로를 정확하게 확인하고, 파일이 존재하는지 확인합니다.

2) RuntimeError: Expected object of type torch.FloatTensor but found type torch.cuda.FloatTensor

GPU에서 저장된 모델을 CPU에서 로드하려고 할 때 발생할 수 있습니다.

• 해결 방법: torch.load() 함수의 map_location 인자를 사용하여 로드할 장치를 지정합니다.

device = torch.device('cpu')
model.load_state_dict(torch.load(model_path, map_location=device))

3) KeyError

state_dict에 예상치 못한 키가 있거나, 저장된 가중치와 모델 구조의 이름이 일치하지 않는 경우 발생합니다.

• 해결 방법: state_dict의 키를 확인하고, 모델 구조를 정확하게 정의했는지 확인합니다. 모델 구조가 변경된 경우, state_dict의 키를 수정하거나, 모델을 다시 학습해야 할 수 있습니다.

4) ModuleNotFoundError 또는 ImportError

모델 정의에 필요한 모듈이나 클래스를 임포트할 수 없는 경우 발생합니다.

• 해결 방법: 필요한 모듈이 설치되어 있는지 확인하고, 모델 정의 파일의 경로를 확인합니다.

5. 결론

PyTorch에서 모델 저장 및 로드는 딥러닝 프로젝트의 핵심적인 부분입니다. 모델의 재사용, 배포, 실험 관리, 학습 중단 시점부터의 재개 등 다양한 이점을 제공합니다. 본 문서에서 설명한 모델 저장 및 로드 방법을 이해하고, 각 상황에 맞는 방법을 선택하여 딥러닝 프로젝트의 효율성을 높일 수 있습니다. 또한, 발생 가능한 문제점들을 미리 인지하고, 해결 방법을 숙지하여 모델 저장 및 로드 과정에서 발생할 수 있는 오류를 최소화할 수 있습니다. 모델 저장 및 로드는 단순히 파일을 저장하고 불러오는 것을 넘어, 딥러닝 모델의 생명 주기를 관리하고, 프로젝트의 성공적인 완성을 위한 중요한 단계입니다.