Tak's Notebook

Kaggle, Machine Learning, Engineering

joblib と pickle の比較

これは何か?

  • pickle でファイル保存できるのとほぼ同じ書き方で保存できる joblib.dump/joblib.load について調べた。
  • pickle と joblib の比較について。compress を使うと joblib の方がファイルサイズが小さくなる。読み込み速度も変わるかもしれない。
  • joblib の compress の値は3程度が良いと言われているが、実際調べてると大体3くらいでちょうどいいことが確認できた。

背景

pickle 以外にファイル保存する方法について調べていて、 joblib がほぼ同じ書き方で圧縮保存できて便利そうだということで興味を持った。
そこで、メモ代わりに joblib について書いてあるブログ記事を twitter でシェアしたところ、意外と反応が(自分の中では)大きかったので、色々調べてまとめてみたという話です。

詳しくはドキュメントをどうぞ。

joblib.dump — joblib 0.13.1 documentation

pickle と joblib の比較

以下、jupyter 上で実行していると思ってください。

import os
import sys
import pickle
import joblib
import numpy as np
from keras.datasets import cifar10

(x_train, y_train), (x_test, y_test) = cifar10.load_data()
x_train = np.array(x_train, dtype=float)  # ファイルサイズを大きくするため無駄に型変換(int->float)
x_train_size = sys.getsizeof(eval("x_train")) / (1024*1024)
print("memory size of X: {} Mb".format(round(x_train_size,1)))
memory size of X: 1171.9 Mb

まず、適当なある程度の大きさのファイルを用意したかったので、Keras にある cifar10 のデータを持ってきて、データ型を int から float に変換しました。ファイルサイズが見事(?) 1GB 程度になったので、これを用いてオブジェクトの保存等を行っていきます。

保存時間

pickle

%%time
with open("x_train.pickle", mode="wb") as f:
    pickle.dump(x_train, f)
CPU times: user 1.9 s, sys: 5.57 s, total: 7.46 s
Wall time: 13.3 s

joblib

%%time
with open("x_train.joblib", mode="wb") as f:
    joblib.dump(x_train, f, compress=3)
CPU times: user 22 s, sys: 1.2 s, total: 23.2 s
Wall time: 24.1 s

両者を比較すると圧縮している作業分、joblib.dumpの方が実行時間がかかっています。

ファイルサイズ

def getsize_filepath(path, filename):
    print( "{}: {} Mb".format(filename, round(os.path.getsize(path)/(1024*1024),1)) )

getsize_filepath("x_train.pickle", "pickle")
getsize_filepath("x_train.joblib", "joblib")
pickle: 1171.9 Mb
joblib: 225.0 Mb

compress が効いているため、このケースだと4倍以下にまで圧縮されています。

読み込み時間

pickle

%%time
with open("x_train.pickle", mode="rb") as f:
    x_train = pickle.load(f)
CPU times: user 469 ms, sys: 3.5 s, total: 3.97 s
Wall time: 6.93 s

joblib

%%time
with open("x_train.joblib", mode="rb") as f:
    x_train = joblib.load(f)
CPU times: user 3.42 s, sys: 619 ms, total: 4.04 s
Wall time: 4.36 s

ファイルサイズが小さくなると読み込み時間に差が出るのかはいまいち検証できませんでした。何回か実行してみると、それぞれ時間にバラツキがあって pickle の方が読み込み早かったりするケースもあったので。

joblib の compress について

pickle との比較の結果 joblib も悪くないなと思うに至りました。次に疑問に思ったのが、 compress の値をいくつにするかということです。
公式によると0~9の値が取りうる値で、3が推奨されているようです。(他にも compress の指定の仕方はあるが割愛)

そこで、 compress の値を変えると、dump の時間・ファイルサイズ・load の時間がどのように変化するかを調べてみました。

結果は以下のようになりました。

f:id:takaishikawa42:20190117224247p:plain
joblib の compress の値に対する dump time, file size, load time の比較

この結果を見ると、compress の値に 3 を推奨するのは納得という印象を受けました。今回のケースを見る限りでは、その値は大きくても 5 くらいが限界かなという感じです。(当然、個々のケースによって異なるかと思いますが)

まとめ

joblib は pickle と比較しても悪くなさそうで、joblib.dumpcompress の値は 3(推奨値)から 5 くらいが良さそうということがわかった。


今回分からなかったこと

  • dump/load 時のメモリの比較
  • 保存ファイルの拡張子名(pickle は .pickle とか .pkl とか使われているけど、joblib はどうするのが良いか。今回は判別しやすいので .joblib としたが、どういった名前にするのが良いか)
  • pandas の feather との比較



今回使ったコードはこちらに置いています。

notebooks/pickle_joblib.ipynb at master · takaiyuk/notebooks · GitHub

notebooks/joblib_compress.ipynb at master · takaiyuk/notebooks · GitHub