[MLOps] 학습이 끝난 ML 모델을 저장하는 방법

1. [머신러닝 모델] Scikit-learn과 같은 라이브러리로 만든 모델을 저장할 경우

1) pickle

- python의 표준 라이브러리로, 별도 설치 없이 사용 가능하며, 간단하다.

- 모든 python 객체(리스트, 딕셔너리, 클래스 인스턴스, 함수 등)을 저장하고 불러올 수 있다.

import pickle
from sklearn.ensemble import RandomForestClassifier

# 모델 생성 및 학습
model = RandomForestClassifier(n_estimators=100)
model.fit(X, y)
# 모델 저장
with open('model.pkl', 'wb') as f:
    pickle.dump(model, f)

# 모델 불러오기
with open('model.pkl', 'rb') as f:
    loaded_model = pickle.load(f)
    
loaded_model.predict([[-1.39710,3.31910,-1.392700,-1.99480 ]])

2) joblib

- 가장 자주 사용되는 방법으로, 특히 Scikit-learn 모델을 저장하는데 적합하다.

- pickle 기반이지만 대규모 데이터를 더 효율적으로 처리한다.

import joblib

# 모델 저장
joblib.dump(model, 'model.joblib')

# 모델 불러오기
loaded_model = joblib.load('model.joblib')

 

(주의)

- pickle과 joblib은 은 데이터를 로드하면서 임의의 python 코드를 실행할 수 있으므로, 신뢰할 수 없는 pickle파일은 열면 안된다.

(대안)

- 모델 자체가 아니라, 메타데이터나 간단한 구조를 저장할때는 JSON 구조 사용

 

2. [딥러닝] TensorFlow, PyTorch 등의 프레임워크를 사용하는 경우

1) TensorFlow, Keras: HDF5 또는 savedModel사용

from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Dense

# 모델 생성
model = Sequential([
    Dense(10, input_shape=(5,), activation='relu'),
    Dense(1, activation='sigmoid')
])

#### 방법1 HDF5로 저장
model.save('model.h5')

# 모델 불러오기
from tensorflow.keras.models import load_model
loaded_model = load_model('model.h5')

#### 방법2 SavedModel로 저장
model.save('model_saved')

# SavedModel 불러오기
loaded_model = load_model('model_saved')

 

2) PyTorch: .pth파일

- Pytorch 모델을 저장할 때는 전체 모델 또는 가중치를 저장하는 방식이 일반적이다.

import torch
import torch.nn as nn

# 간단한 모델 생성
model = nn.Linear(10, 1)

#### 방법1. 가중치저장
torch.save(model.state_dict(), 'model_weights.pth')

# 가중치 불러오기
model.load_state_dict(torch.load('model_weights.pth'))
model.eval()

#### 방법2. 전체 모델 저장
torch.save(model, 'model.pth')

# 전체 모델 불러오기
model = torch.load('model.pth')
model.eval()