pickle / joblib 를 사용한 머신러닝 모델 저장하는 법

대학시절 pickle 모듈을 사용했었는데, pickle모듈을 이용해 저장하려면 먼저 파일을 바이너리 모드로 열고 저장해야하므로 아무래도 조금 번거로운 감이 있었다. 인터넷 서치를 해보니, joblib 모듈을 이용하면 위 예시처럼 dump함수와 load함수 만으로 모델을 바로 저장하고 불러올 수 있어 훨씬 직관적이다.

 

학습시킨 모델 저장하기

다음은 K최근접 이웃 알고리즘을 이용해 붓꽃 품종 분류 모델을 학습시키고 
joblib 모듈을 이용해 pickle 파일로 저장하는 코드이다.

from sklearn import datasets
from sklearn.neighbors import KNeighborsClassifier
import joblib

# 붓꽃 품종 데이터 세트 불러오기
dataset = datasets.load_iris()

# 입력 데이터와 타깃을 준비합니다.
X, y = dataset['data'], dataset['target']

# K 최근접이웃 모델 객체를 만듭니다.
model = KNeighborsClassifier(n_neighbors=5)

# K 최근접 이웃 모델에 입력 데이터와 타깃을 입력하고 학습시킨다.
model.fit(X,y)

# 학습시킨 모델을 현재 경로에 knn_model.pkl 파일로 저장합니다.
joblib.dump(model, './knn_model.pkl')

저장한 모델을 불러와 사용하기

다음은 pickle 파일로 저장된 모델을 불러와 정확도를 측정하는 코드이다.

import joblib

loaded_model = joblib.load('./knn_model.pkl')

score = loaded_model.score(X,y)
print('정확도: {score:.3f}'.format(score=score))