Python: scikit-learn の Pipeline 機能のキャッシュを試す

今回は scikit-learn の Pipeline に存在するキャッシュの機能を試してみる。 scikit-learn の Pipeline は、データセットの前処理・特徴量抽出からモデルの学習・推論までの一連の処理手順をひとまとめにして扱うことのできる機能。以前に、このブログでも扱ったことがある。

blog.amedama.jp

機械学習のタスクにおいて、データセットの前処理や特徴量抽出には意外と時間がかかる。大抵の場合、特徴量エンジニアリングを通してモデルで使う特徴量はどんどん増えていく。そうした状況で、同じ特徴量を毎回ゼロから計算するのは時間がもったいない。そこで、計算済みの特徴量についてはキャッシュしておきたくなる。そのニーズを Pipeline で満たすのがキャッシュ機能になる。

今回使った環境は次の通り。

$ sw_vers
ProductName:    Mac OS X
ProductVersion: 10.13.6
BuildVersion:   17G65
$ python -V
Python 3.7.0

下準備

まずは scikit-learn をインストールしておく。

$ pip install scikit-learn

キャッシュ機能を試す

早速だけど、以下が Pipeline でキャッシュを有効にしたサンプルコードになる。今回の本質ではないけど、この中では Digits データセットを主成分分析して、それをランダムフォレストで識別している。ポイントは Pipeline クラスのコンストラクタで memory オプションに sklearn.externals.joblib.memory.Memory オブジェクトを渡しているところ。これで Pipeline のキャッシュが有効になる。また、キャッシュの効果を確認するために学習と推論にかかった時間を計測している。

#!/usr/bin/env python
# -*- coding: utf-8 -*-

from contextlib import contextmanager
import time

from sklearn import datasets
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import accuracy_score
from sklearn.pipeline import Pipeline
from sklearn.decomposition import KernelPCA
from sklearn.model_selection import train_test_split
from sklearn.externals.joblib import memory


@contextmanager
def elapsed(section):
    """コードブロックの実行時間を測るためのコンテキストマネージャ"""
    start_time = time.time()
    yield
    end_time = time.time()
    print(f'{section} elapsed: {end_time - start_time} sec')


def main():
    # Digits データセットを読み込む
    dataset = datasets.load_digits()
    X, y = dataset.data, dataset.target

    # ホールドアウト検証用にデータを分割する
    X_train, X_test, y_train, y_test = train_test_split(X, y, shuffle=True, test_size=0.33, random_state=42)

    # 処理のパイプライン
    steps = [
        # 主成分分析
        ('pca', KernelPCA()),
        # ランダムフォレスト
        ('rf', RandomForestClassifier(n_estimators=100))
    ]

    # pipeline-cache ディレクトリにキャッシュデータを保存する
    cache = memory.Memory(location='pipeline-cache', verbose=3)
    pipeline = Pipeline(steps=steps, memory=cache)

    # 学習フェーズにかかる時間を計測する
    with elapsed('fit'):
        pipeline.fit(X_train, y_train)

    # 推論フェーズにかかる時間を計測する
    with elapsed('pred'):
        y_pred = pipeline.predict(X_test)
        acc = accuracy_score(y_pred, y_test)
        print(f'accuracy: {acc}')


if __name__ == '__main__':
    main()

では、上記のサンプルコードを実行してみよう。最初の実行では、当然ながらキャッシュは効かない。

$ python pipecache.py 
________________________________________________________________________________
[Memory] Calling sklearn.pipeline._fit_transform_one...
_fit_transform_one(KernelPCA(alpha=1.0, coef0=1, copy_X=True, degree=3, eigen_solver='auto',
     fit_inverse_transform=False, gamma=None, kernel='linear',
     kernel_params=None, max_iter=None, n_components=None, n_jobs=None,
     random_state=None, remove_zero_eig=False, tol=0), 
array([[0., ..., 0.],
       ...,
       [0., ..., 0.]]), array([1, ..., 1]), None)
________________________________________________fit_transform_one - 0.4s, 0.0min
fit elapsed: 1.4524390697479248 sec
accuracy: 0.9478114478114478
pred elapsed: 0.05564618110656738 sec

学習には 1.45 秒かかった。尚、これには前処理・特徴量抽出の時間も含まれている。今回に関して言えばデータセットを主成分分析してランダムフォレストを学習させるのにかかった時間ということになる。

それでは、もう一度同じように実行してみよう。二回目はキャッシュが効くので、出力も少し異なったものになる。

$ python pipecache.py
[Memory]0.0s, 0.0min    : Loading _fit_transform_one...
fit elapsed: 1.1166560649871826 sec
accuracy: 0.9562289562289562
pred elapsed: 0.055490970611572266 sec

今度は学習が 1.11 秒と、先ほどよりも短い時間で終わっている。これがキャッシュの効果といえる。

ちなみに、キャッシュされるものは Pipeline の最終段に位置するモデル (Estimator) に渡される直前のデータになる。つまり、前処理や特徴量抽出の部分がキャッシュされるのであって、モデルの学習結果自体はキャッシュされない。これは、以下のソースコードを読むと確認できる。

github.com

もし、学習結果を含めて保存したいなら、学習済みの Pipeline オブジェクトそのものを Pickle で保存しておくのが良い。

blog.amedama.jp

ちなみに、キャッシュに使う joblib の Memory は、やろうと思えば自分で拡張もできそうな雰囲気がある。

github.com

例えば一次キャッシュをローカルのディスクにして、二次キャッシュをクラウド上のオブジェクトストレージにしたような実装もできそうだ。そうしたものがあれば、どんな環境からでもキャッシュ済みのデータが利用できて時間の節約になりそうに感じた。 GCP の Preemptible VM なんか使ってるときは特に便利そう。