모델 저장 및 로드 방법 (TensorFlow & PyTorch)

모델 저장 및 로드 방법 (TensorFlow & PyTorch)

딥러닝 모델을 학습하고 난 후 이를 저장하고 불러오는 과정은 매우 중요합니다. 학습된 모델을 재사용하거나 배포할 때 필요한 과정이며, 훈련을 중단한 후에도 모델을 다시 불러와 이어서 학습할 수 있습니다. 본 포스팅에서는 TensorFlow와 PyTorch에서 모델을 저장하고 로드하는 방법을 설명하겠습니다.

1. TensorFlow에서 모델 저장 및 로드

TensorFlow에서는 tf.keras.Model을 이용하여 모델을 저장할 수 있습니다. TensorFlow에서 제공하는 모델 저장 방법은 크게 두 가지로 나뉩니다.

1.1 전체 모델 저장 및 로드

TensorFlow에서는 model.save()를 사용하여 전체 모델을 저장할 수 있습니다. 이 방법은 모델의 가중치뿐만 아니라 모델의 구조, 옵티마이저 설정, 컴파일 정보까지 포함합니다.

모델 저장하기

import tensorflow as tf
from tensorflow import keras

# 간단한 모델 정의
model = keras.Sequential([
    keras.layers.Dense(64, activation='relu', input_shape=(10,)),
    keras.layers.Dense(1, activation='sigmoid')
])

model.compile(optimizer='adam', loss='binary_crossentropy', metrics=['accuracy'])

# 모델 저장
model.save("saved_model/my_model")

모델 불러오기

저장된 모델을 다시 불러올 때는 tf.keras.models.load_model()을 사용하면 됩니다.

# 저장된 모델 불러오기
new_model = keras.models.load_model("saved_model/my_model")

# 불러온 모델 요약
new_model.summary()

1.2 가중치만 저장 및 로드

모델의 구조는 그대로 두고 가중치만 저장하는 방법도 있습니다. 이는 model.save_weights()와 model.load_weights()를 사용합니다.

가중치 저장하기

# 가중치 저장
model.save_weights("weights/my_weights.h5")

가중치 불러오기

가중치를 불러오려면 동일한 모델 구조를 먼저 정의한 후, load_weights()를 호출해야 합니다.

# 동일한 모델 정의
model = keras.Sequential([
    keras.layers.Dense(64, activation='relu', input_shape=(10,)),
    keras.layers.Dense(1, activation='sigmoid')
])

model.compile(optimizer='adam', loss='binary_crossentropy', metrics=['accuracy'])

# 가중치 불러오기
model.load_weights("weights/my_weights.h5")

---

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

PyTorch에서는 torch.save()와 torch.load()를 사용하여 모델을 저장하고 불러올 수 있습니다. PyTorch에서는 주로 다음 두 가지 방법이 있습니다.

2.1 전체 모델 저장 및 로드

PyTorch에서는 torch.save(model, path)를 이용하여 전체 모델을 저장할 수 있습니다.

모델 저장하기

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

# 간단한 모델 정의
class SimpleModel(nn.Module):
    def __init__(self):
        super(SimpleModel, self).__init__()
        self.fc1 = nn.Linear(10, 64)
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(64, 1)
        self.sigmoid = nn.Sigmoid()

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

# 모델 생성
model = SimpleModel()

# 모델 저장
torch.save(model, "pytorch_model.pth")

모델 불러오기

전체 모델을 저장한 경우, 불러올 때도 같은 방법을 사용합니다.

# 저장된 모델 불러오기
loaded_model = torch.load("pytorch_model.pth")

# 모델을 평가 모드로 설정
loaded_model.eval()

2.2 가중치만 저장 및 로드

모델 전체를 저장하는 것보다 가중치만 저장하는 것이 일반적으로 더 권장됩니다. 가중치만 저장하면 모델 클래스 정의가 필요하지만, 더 유연하게 사용할 수 있습니다.

가중치 저장하기

# 모델의 state_dict 저장
torch.save(model.state_dict(), "pytorch_weights.pth")

가중치 불러오기

불러올 때는 동일한 모델 구조를 정의한 후, load_state_dict()를 사용해야 합니다.

# 동일한 모델 정의
model = SimpleModel()

# 가중치 불러오기
model.load_state_dict(torch.load("pytorch_weights.pth"))

# 모델을 평가 모드로 설정
model.eval()

---

3. 저장 방식 비교 및 권장 방법

방법	TensorFlow	PyTorch
전체 모델 저장	model.save("path")	torch.save(model, "path")
전체 모델 불러오기	keras.models.load_model("path")	torch.load("path")
가중치만 저장	model.save_weights("path")	torch.save(model.state_dict(), "path")
가중치만 불러오기	model.load_weights("path")	model.load_state_dict(torch.load("path"))

전체 모델을 저장하는 방법은 모델의 구조와 가중치가 함께 저장되므로 빠르게 복구할 수 있는 장점이 있습니다. 하지만 PyTorch에서는 전체 모델을 저장하는 것보다 가중치만 저장하는 것을 추천합니다. 이는 모델 클래스를 다시 정의해야 하는 번거로움이 있지만, 향후 모델을 확장하거나 재사용할 때 더 유연하기 때문입니다.

TensorFlow의 경우 SavedModel 형식으로 저장하는 것이 가장 표준적인 방법이며, 배포할 때도 쉽게 사용할 수 있습니다.

---

4. 결론

이번 포스팅에서는 TensorFlow와 PyTorch에서 모델을 저장하고 불러오는 방법을 설명하였습니다.

• TensorFlow에서는 model.save()를 이용하여 전체 모델을 저장하거나 save_weights()를 이용하여 가중치만 저장할 수 있습니다.
• PyTorch에서는 torch.save(model, path)를 이용하여 전체 모델을 저장하거나 torch.save(model.state_dict(), path)를 이용하여 가중치만 저장할 수 있습니다.
• PyTorch에서는 가중치만 저장하고 로드하는 방식을 추천합니다.

모델을 저장하고 불러오는 방법을 익히면, 딥러닝 모델을 쉽게 재사용하고 배포할 수 있습니다. 프로젝트에 맞는 방법을 선택하여 효과적으로 활용해 보시길 바랍니다.