[PYTHON] 모델 가중치 저장 및 불러오기 7가지 방법과 직렬화 문제 해결 전략

모델 가중치 저장 및 불러오기

 

딥러닝 모델 개발 과정에서 가장 허망한 순간은 수십 시간 동안 학습시킨 모델의 결과물인 가중치(Weights)를 제대로 저장하지 않아 유실하거나, 프레임워크 버전 차이로 인해 불러오기에 실패하는 경우입니다. 모델의 가중치는 신경망이 데이터로부터 습득한 '지능의 정수'이며, 이를 안전하게 직렬화(Serialization)하고 복원하는 기술은 배포 및 유지보수의 핵심입니다. 본 가이드에서는 파이썬 기반의 PyTorch, TensorFlow, Scikit-learn 등 주요 프레임워크에서 가중치를 관리하는 7가지 전문적인 방법과 실무에서 발생하는 호환성 해결책을 심도 있게 다룹니다.

---

1. 가중치 저장의 내부 메커니즘: state_dict vs Entire Model

모델을 저장할 때는 단순히 파일로 만드는 것을 넘어, 어떤 정보를 담을지 결정해야 합니다. 가중치만 담긴 상태 사전(state_dict)을 저장할 것인지, 모델의 구조(Architecture)까지 포함한 전체 객체를 저장할 것인지에 따라 활용도가 달라집니다.

비교 항목	가중치만 저장 (Weight Only / state_dict)	전체 모델 저장 (Entire Model)
저장 내용	학습된 파라미터(Weight, Bias) 값만 포함	모델 구조 정보 + 파라미터 값 모두 포함
유연성	매우 높음 (코드로 구조 정의 후 값만 주입)	낮음 (동일한 클래스 정의가 환경에 있어야 함)
파일 용량	상대적으로 작음	상대적으로 큼
권장 사용처	실무 배포, 미세 조정(Fine-tuning)	빠른 테스트, 학습 중단점 복구(Checkpoint)
해결 방법	구조 코드 공유가 필수적임	피클(Pickle) 관련 보안 및 경로 문제 주의

---

2. 실무자를 위한 모델 가중치 관리 노하우 7가지 (Python Code)

개발자가 실무 현장에서 즉시 복사하여 사용할 수 있는 프레임워크별 표준 구현 사례입니다.

Example 1: PyTorch의 권장 방식 - state_dict 저장 및 로드

모델의 구조는 파이썬 코드로 유지하고 가중치 값만 `.pth` 파일로 관리하는 가장 표준적인 방법입니다.

import torch
import torch.nn as nn

# 1. 모델 정의
model = nn.Sequential(nn.Linear(10, 5), nn.ReLU(), nn.Linear(5, 2))

# 2. 가중치 저장 (권장 확장자: .pth 또는 .pt)
torch.save(model.state_dict(), 'model_weights.pth')

# 3. 불러오기 (먼저 동일한 구조의 모델 객체 생성 필수)
new_model = nn.Sequential(nn.Linear(10, 5), nn.ReLU(), nn.Linear(5, 2))
new_model.load_state_dict(torch.load('model_weights.pth'))
new_model.eval() # 추론 모드로 전환
print("PyTorch 가중치 복원 성공")
    

Example 2: TensorFlow/Keras - H5 및 SavedModel 형식 활용

구글의 생태계에서 모델을 저장하는 두 가지 주요 해결 방법입니다.

import tensorflow as tf

# 모델 생성 및 학습 가정
model = tf.keras.Sequential([tf.keras.layers.Dense(10, input_shape=(5,))])

# 방식 A: HDF5 포맷 (단일 파일)
model.save_weights('weights.h5')
model.load_weights('weights.h5')

# 방식 B: SavedModel 포맷 (디렉토리 구조, 서빙 최적화)
model.save('saved_model_v1')
new_tf_model = tf.keras.models.load_model('saved_model_v1')
    

Example 3: 학습 중단 방지를 위한 체크포인트(Checkpoint) 자동 저장

에폭(Epoch)마다 최상의 성능을 내는 가중치를 자동으로 기록하는 실무 패턴입니다.

from tensorflow.keras.callbacks import ModelCheckpoint

checkpoint_path = "checkpoints/model-{epoch:02d}-{val_loss:.2f}.hdf5"
checkpoint_callback = ModelCheckpoint(
    filepath=checkpoint_path,
    save_weights_only=True,
    monitor='val_loss',
    mode='min',
    save_best_only=True
)
# model.fit(..., callbacks=[checkpoint_callback])
    

Example 4: Scikit-learn 모델의 직렬화 (Joblib 활용)

딥러닝 외에 전통적 머신러닝 모델을 대용량 데이터 환경에서 저장할 때의 차이점 해결책입니다.

import joblib
from sklearn.ensemble import RandomForestClassifier

clf = RandomForestClassifier()
# 모델 저장
joblib.dump(clf, 'rf_model.pkl')

# 모델 로드
loaded_clf = joblib.load('rf_model.pkl')
print("Scikit-learn 모델 로드 완료")
    

Example 5: GPU에서 학습한 모델을 CPU 환경에서 불러오기

배포 환경이 CPU일 때 발생하는 장치 불일치(RuntimeError) 해결 방법입니다.

# PyTorch: map_location 파라미터 활용이 핵심
device = torch.device('cpu')
weights = torch.load('model_weights.pth', map_location=device)
model.load_state_dict(weights)
    

Example 6: 가중치 파일에서 특정 레이어만 추출하여 전이 학습에 적용

전체 모델이 아닌 특정 계층의 지능만 빌려오는 고수준의 해결 전략입니다.

pretrained_dict = torch.load('expert_model.pth')
model_dict = model.state_dict()

# 현재 모델과 가중치 이름이 일치하는 파라미터만 필터링
pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict}
model_dict.update(pretrained_dict)
model.load_state_dict(model_dict)
    

Example 7: ONNX 포맷 변환을 통한 프레임워크 간 가중치 공유

PyTorch로 만든 가중치를 다른 언어나 환경(C++, 런타임 등)에서 쓰기 위한 변환법입니다.

import torch.onnx

dummy_input = torch.randn(1, 10)
torch.onnx.export(model, dummy_input, "model.onnx", 
                  input_names=['input'], output_names=['output'])
print("ONNX 범용 포맷으로 가중치 내보내기 성공")
    

---

3. 가중치 로드 실패 시 트러블슈팅 및 해결 전략

현업에서 가장 자주 발생하는 3가지 오류와 그 해결책을 제시합니다.

1. KeyError (Missing key(s) in state_dict): 모델 구조 코드를 수정했는데 예전 가중치를 불러올 때 발생합니다. `strict=False` 옵션을 주어 일치하는 부분만 로드하게 하세요.
2. Size Mismatch: 입력 데이터의 차원이 달라졌을 때 발생합니다. 가중치 행렬의 쉐이프(Shape)를 확인하고 전처리 단계의 차이를 해결해야 합니다.
3. Version Conflict: 프레임워크 버전이 다를 때 생기는 문제입니다. 가급적 가중치와 함께 사용된 라이브러리의 버전을 `requirements.txt`에 명시하고 관리하세요.

---

4. 결론 및 향후 전망

가중치 저장 기술은 단순히 데이터를 기록하는 행위를 넘어, AI 모델의 재현성(Reproducibility)과 직결됩니다. 최근에는 모델 가중치의 용량이 기하급수적으로 커짐에 따라(LLM 등), 이를 효율적으로 압축하여 저장하는 양자화(Quantization) 저장 방식이 새로운 표준으로 자리 잡고 있습니다. 본 가이드의 7가지 기초 및 심화 기법을 숙지한다면, 어떤 환경에서도 유연하게 모델을 복원하고 서비스에 투입할 수 있을 것입니다.

 

참고 문헌 (Sources):

• PyTorch Tutorials: Saving and Loading Models (pytorch.org)
• TensorFlow Guide: Save and load Keras models (tensorflow.org)
• ONNX (Open Neural Network Exchange) Documentation.
• "Deep Learning with Python" by François Chollet.