CUBE SUGAR CONTAINER

技術系のこと書きます。

Python: LightGBM の cv() 関数から学習済みモデルを得る

勾配ブースティング決定木を扱うフレームワークの一つである LightGBM の Python API には cv() という関数がある。 この "cv" というのは Cross Validation の略で、その名の通り LightGBM のモデルを交差検証するための関数になっている。 具体的には、この関数にデータセットを渡すと、そのデータでモデルを学習させると共に、指定した評価指標について交差検証で評価できる。

今回は、この関数から交差検証の過程で学習させたモデルを手に入れる方法について書いてみる。 というのも、この関数が返すのは指定した評価指標を用いて計測された性能に関する情報に限られているため。 ようするに、交差検証の過程で学習させた学習済みのモデルを手に入れる方法が、標準では用意されていない。 しかしながら、交差検証で性能が確かめられたモデルを取得したい、というニーズはあるはず。

なお、結果から先に書くとコールバック関数を使うことで学習済みモデルを手に入れることができた。

使った環境は次の通り。

$ sw_vers
ProductName:    Mac OS X
ProductVersion: 10.14.4
BuildVersion:   18E226
$ python -V          
Python 3.7.3
$ pip list | grep -i lightgbm
lightgbm        2.2.3  

下準備

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

$ brew install libomp
$ pip install lightgbm scikit-learn

cv() 関数からコールバック関数を使って学習済みモデルを取り出す

論よりコードということで、以下に cv() 関数からコールバック関数を使って学習済みモデルを取り出すサンプルコードを示す。 LightGBM の train() 関数や cv() 関数には、ブースティングのイテレーションごとに呼ばれるコールバック関数が登録できる。 一般的には、コールバック関数は学習の過程などを記録するために用いられる。 しかしながら、コールバック関数にはモデルも渡されるため、今回のような用途にも応用が効く。 内容の細かい解説に関してはコメントの形でコードに含めた。

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

import lightgbm as lgb
from sklearn import datasets
from sklearn.metrics import accuracy_score
from sklearn.model_selection import train_test_split
import numpy as np


class ModelExtractionCallback(object):
    """lightgbm.cv() から学習済みモデルを取り出すためのコールバックに使うクラス

    NOTE: 非公開クラス '_CVBooster' に依存しているため将来的に動かなく恐れがある
    """

    def __init__(self):
        self._model = None

    def __call__(self, env):
        # _CVBooster の参照を保持する
        self._model = env.model

    def _assert_called_cb(self):
        if self._model is None:
            # コールバックが呼ばれていないときは例外にする
            raise RuntimeError('callback has not called yet')

    @property
    def boosters_proxy(self):
        self._assert_called_cb()
        # Booster へのプロキシオブジェクトを返す
        return self._model

    @property
    def raw_boosters(self):
        self._assert_called_cb()
        # Booster のリストを返す
        return self._model.boosters

    @property
    def best_iteration(self):
        self._assert_called_cb()
        # Early stop したときの boosting round を返す
        return self._model.best_iteration


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

    # デモ用にデータセットを分割する
    X_train, X_test, y_train, y_test = train_test_split(X, y,
                                                        shuffle=True,
                                                        random_state=42)

    # LightGBM 用のデータセット表現に直す
    lgb_train = lgb.Dataset(X_train, y_train)

    # 学習済みモデルを取り出すためのコールバックを用意する
    extraction_cb = ModelExtractionCallback()
    callbacks = [
        extraction_cb,
    ]

    # データセットを 5-Fold CV で学習する
    lgbm_params = {
        'objective': 'multiclass',
        'num_class': 3,
    }
    # NOTE: 一般的には返り値の内容 (交差検証の結果) を確認する
    lgb.cv(lgbm_params,
           lgb_train,
           num_boost_round=1000,
           early_stopping_rounds=10,
           nfold=5,
           shuffle=True,
           stratified=True,
           seed=42,
           callbacks=callbacks,
           )

    # コールバックのオブジェクトから学習済みモデルを取り出す
    proxy = extraction_cb.boosters_proxy
    boosters = extraction_cb.raw_boosters
    best_iteration = extraction_cb.best_iteration

    # 各モデルの推論結果を Averaging する場合
    y_pred_proba_list = proxy.predict(X_test,
                                      num_iteration=best_iteration)
    y_pred_proba_avg = np.array(y_pred_proba_list).mean(axis=0)
    y_pred = np.argmax(y_pred_proba_avg, axis=1)
    accuracy = accuracy_score(y_test, y_pred)
    print('Averaging accuracy:', accuracy)

    # 各モデルで個別に推論する場合
    for i, booster in enumerate(boosters):
        y_pred_proba = booster.predict(X_test,
                                       num_iteration=best_iteration)
        y_pred = np.argmax(y_pred_proba, axis=1)
        accuracy = accuracy_score(y_test, y_pred)
        print('Model {0} accuracy: {1}'.format(i, accuracy))

if __name__ == '__main__':
    main()

上記の実行結果は次の通り。 5-Fold CV から得られた全モデルを Averaging する場合と、各モデルごとに推論させた場合の精度 (Accuracy) が示される。

$ python lgbcv.py
...(snip)...
Averaging accuracy: 1.0
Model 0 accuracy: 1.0
Model 1 accuracy: 1.0
Model 2 accuracy: 0.9473684210526315
Model 3 accuracy: 1.0
Model 4 accuracy: 1.0

なお、今回のコードで注意すべきなのは、非公開のクラスを利用しているところがある点。 具体的には ModelExtractionCallback#boosters_proxy から得られるオブジェクトが lightgbm.engine._CVBooster というクラスのインスタンスになっている。 このクラスは、名前の先頭にアンダースコア (_) が含まれることから、外部に API として公開しているものではないと考えられる。 そのため、このモデルのインターフェースが唐突に変わったとしても文句はいえない。

一応、今の _CVBooster がどういったコードになっているのか示しておく。 実装があるのは以下の場所。 このクラスは内部にモデルをリストの形で保持している。 そして、オブジェクトに何らかの呼び出しがあると、特殊メソッドの __getattr__() でトラップされる。 トラップされた呼び出しは、内部のモデルに対して順番に実行されて結果がリストとして返されることになる。

github.com

いじょう。

(2019-04-15 追記)

この記事に掲載したコードをベースに OOF Prediction の実装を追加した Kaggle カーネルをまつけんさん (Twitter: @Kenmatsu4) が公開してくださいました。 ありがとうございます。

www.kaggle.com