pickle と joblib で機械学習モデルを保存＆読み込む（Python）

scikit-learnなどで構築した機械学習モデル（オブジェクト）を、どのようにファイルに保存し、それを読み込み使うのか。

よくあるのが以下の2つです。

pickle：保存＆読み込むスピードがjoblibより早い
joblib：保存されるファイルサイズがpickleより小さい

違いは、スピードかサイズか、ということになります。

今回は、「pickle と joblib で機械学習モデルを保存＆読み込む（Python）」というお話しと言うか、実行例です。

Contents

今回の流れ

今回の流れです。

ライブラリー読み込み
サンプルデータ生成（回帰問題）
データ分割（学習データとテストデータ）
学習データを使ってモデル生成（ランダムフォレスト）
pickleを使ったモデル保存・読み込み・予測
joblibを使ったモデル保存・読み込み・予測

最後の2つのところで、以下の計測をします。

保存スピード
保存ファイルのサイズ
読み込みスピード
テストデータの予測精度（R2）

最後の予測精度（R2）ですが、pickle と joblib とも同じモデルを保存した後に読み込んでいるので、同じ値になります。

ライブラリー読み込みからモデル生成まで

以下、コードです。

# 1. ライブラリー読み込み
import pandas as pd
from sklearn.datasets import make_regression
from sklearn.model_selection import train_test_split
from sklearn.ensemble import RandomForestRegressor
from sklearn.metrics import r2_score
import pickle
import joblib
import os
import time

# 2. サンプルデータ生成（回帰問題）
X, y = make_regression(n_samples=10000,   #サンプルの数
                       n_features=1000,   #特徴量の数
                       noise=10,          #ノイズの標準偏差
                       bias=0,            #バイアス（切片）
                       random_state=12)

X = pd.DataFrame(X) #説明変数（特徴量）
y = pd.DataFrame(y) #目的変数

# 3. データ分割（学習データとテストデータ）
X_train, X_test, y_train, y_test = train_test_split(X,
                                                    y,
                                                    train_size=0.7,
                                                    test_size=0.3,
                                                    random_state=42)

# 4. 学習データを使ってモデル生成（ランダムフォレスト）
model =RandomForestRegressor(n_jobs=-1).fit(X_train,y_train)

今回は、回帰問題用のサンプルデータを生成した後に、このサンプルデータをモデル構築用の学習データと構築したモデルを検証するテストデータに分割し、学習データを使いランダムフォレストで回帰モデルを構築しました。

pickleを使ったモデル保存・読み込み・予測

先ず、pickleで構築したモデルを保存します。

以下、コードです。

%%time
# 構築したモデルの保存
filename = 'model.pkl'
pickle.dump(model,open(filename,'wb'))

以下、実行結果です。

保存するのに、60.2 ms（ミリ秒）かかっていることが分かります。

この保存したファイルのサイズを確かめてみます。

以下、コードです。

# 保存したモデルのサイズ
print(round(os.path.getsize(filename)/1024/1024,1),'Mb')

以下、実行結果です。

ファイルサイズは、54.1 Mbです。

次に、今保存したファイルを読み込みます。

以下、コードです。

%%time
# 保存したモデルの読み込み
loaded_rf_model = pickle.load(open(filename, 'rb'))

以下、実行結果です。

読み込むのに、71.1 ms（ミリ秒）かかっていることが分かります。

この読み込んだファイル（予測モデル）で、テストデータを使い予測し精度評価（R2）します。

以下、コードです。

# テストデータによる精度検証（R2）
y_pred = loaded_rf_model.predict(X_test)
r2_score(y_test, y_pred)

以下、実行結果です。

R2（決定係数）は、84.54%です。

joblibを使ったモデル保存・読み込み・予測

先ず、joblibで構築したモデルを保存します。

以下、コードです。

%%time
# 構築したモデルの保存
filename = 'model.pkl2'
joblib.dump(model,filename, compress=3)

以下、実行結果です。

保存するのに、783 ms（ミリ秒）かかっていることが分かります。

pickleに比べると、倍以上時間が掛かっています。これは主にファイル圧縮の処理が入っているためです。

この保存したファイルのサイズを確かめてみます。

以下、コードです。

# 保存したモデルのサイズ
print(round(os.path.getsize(filename)/1024/1024,1),'Mb')

以下、実行結果です。

ファイルサイズは、17.6 Mbです。

pickleに比べると、ファイルサイズが半分以下になっていることが分かります。

次に、今保存したファイルを読み込みます。

以下、コードです。

%%time
# 保存したモデルの読み込み
loaded_rf_model = joblib.load(filename)

以下、実行結果です。

読み込むのに、271 ms（ミリ秒）かかっていることが分かります。

この読み込んだファイル（予測モデル）で、テストデータを使い予測し精度評価（R2）します。

以下、コードです。

# テストデータによる精度検証（R2）
y_pred = loaded_rf_model.predict(X_test)
r2_score(y_test, y_pred)

以下、実行結果です。

R2（決定係数）は、84.54%です。

まとめ

今回は、今回は、「pickle と joblib で機械学習モデルを保存＆読み込む（Python）」というお話しと言うか、実行例を示しました。

この2つの主な違いは、保存＆読み込むスピードと、保存されたファイルサイズにあります。

pickle：保存＆読み込むスピードがjoblibより早い
joblib：保存されるファイルサイズがpickleより小さい

参考にして頂ければと思います。