Python: Optuna の LightGBMTunerCV から学習済みモデルを取り出す

Optuna v1.5.0 では、LightGBM インテグレーションの一環として LightGBMTunerCV という API が追加された。これは LightGBM の cv() 関数を Step-wise algorithm で最適化するラッパーになっている。つまり、重要ないくつかのパラメータを Step-wise で調整することで、最も高い交差検証スコアが得られるパラメータを探索できる。今回は、追加された LightGBMTunerCV の使い方を紹介すると共に学習済みモデルを取り出す方法について書いてみる。ただし、今のところ Experimental な機能という位置づけなので、今後インターフェースなどが変わる可能性もある。

尚、LightGBM の素の cv() 関数から学習済みモデルを取り出す方法は次のエントリに書いた。

blog.amedama.jp

使った環境は次のとおり。

$ sw_vers 
ProductName:    Mac OS X
ProductVersion: 10.14.6
BuildVersion:   18G5033
$ python -V
Python 3.7.7
$ pip list | egrep -i "(optuna|lightgbm)"
lightgbm        2.3.1
optuna          1.5.0

下準備

あらかじめ、使用するパッケージをインストールしておく。

$ brew install libomp
$ pip install lightgbm optuna scikit-learn

乳がんデータセットを交差検証するサンプルコード

早速だけど、以下に LightGBMTunerCV の基本的な使い方を紹介するサンプルコードを示す。これまでと同様、optuna.integration.lightgbm パッケージをインポートすることで Optuna 経由で LightGBM を使うことになる。 LightGBMTunerCV はクラスとして実装されていて、インスタンス化するときに cv() 関数を実行するときに使うオプションを渡す。そして、LightGBMTunerCV#run() メソッドを実行すると重要なパラメータが順番に探索されていって、終了するとインスタンスから最も優れたスコアやパラメータが得られる。

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

from pprint import pprint

from optuna.integration import lightgbm as lgb
from sklearn import datasets
from sklearn.model_selection import StratifiedKFold


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

    # LightGBM のデータセット表現にする
    lgb_train = lgb.Dataset(X, y)

    # データセットの分割方法
    folds = StratifiedKFold(n_splits=3, shuffle=True, random_state=42)

    # 最適化するときの条件
    lgbm_params = {
        'objective': 'binary',
        'metric': 'binary_logloss',
        'verbosity': -1,
    }
    # 基本的には cv() 関数のオプションがそのまま渡せる
    tuner_cv = lgb.LightGBMTunerCV(
        lgbm_params, lgb_train,
        num_boost_round=1000,
        early_stopping_rounds=100,
        verbose_eval=20,
        folds=folds,
    )

    # 最適なパラメータを探索する
    tuner_cv.run()

    # 最も良かったスコアとパラメータを書き出す
    print(f'Best score: {tuner_cv.best_score}')
    print('Best params:')
    pprint(tuner_cv.best_params)


if __name__ == '__main__':
    main()

それでは、上記をファイルに保存して実行してみよう。

$ python tunercv.py
/Users/amedama/.virtualenvs/py37/lib/python3.7/site-packages/optuna/_experimental.py:83: ExperimentalWarning: LightGBMTunerCV is experimental (supported from v1.5.0). The interface can change in the future.
  ExperimentalWarning,
feature_fraction, val_score: inf:   0%|                   | 0/7 [00:00<?, ?it/s][20] cv_agg's binary_logloss: 0.194777 + 0.0385719
[40]  cv_agg's binary_logloss: 0.141999 + 0.056194
[60]   cv_agg's binary_logloss: 0.140204 + 0.0759544

...

min_data_in_leaf, val_score: 0.094602: 100%|######| 5/5 [00:01<00:00,  3.55it/s]
Best score: 0.09460200054481203
Best params:
{'bagging_fraction': 0.5636995994183087,
 'bagging_freq': 4,
 'feature_fraction': 0.44800000000000006,
 'lambda_l1': 0.015234706690217153,
 'lambda_l2': 2.2227989818062668e-07,
 'metric': 'binary_logloss',
 'min_child_samples': 20,
 'num_leaves': 3,
 'objective': 'binary',
 'verbosity': -1}

探索した中で、最も優れた交差検証のスコアを出したパラメータが確認できた。

LightGBMTunerCV から学習済みモデルを取り出す

基本的な使い方が分かったところで今回の本題に入る。別に LightGBM に限った話ではないけど、交差検証するときに学習させたモデル群を使って Averaging などをするのは一般的な手法となっている。とはいえ、交差検証をして得られたパラメータを使って、もう一度同じデータを学習させるのは効率が悪い。そこで、LightGBMTunerCV を使ってパラメータを探索すると同時に、最適なパラメータで学習したモデルを取り出してみた。

以下がそのサンプルコードになる。学習済みモデルの取り出しは TunerCVCheckpointCallback というコールバックで実装している。探索の対象となったパラメータごとにモデルへの参照をインスタンスに保持しておいて、後から取り出せるように作った。

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

import numpy as np
from optuna.integration import lightgbm as lgb
from sklearn import datasets
from sklearn.metrics import accuracy_score
from sklearn.model_selection import StratifiedKFold, train_test_split


class TunerCVCheckpointCallback(object):
    """Optuna の LightGBMTunerCV から学習済みモデルを取り出すためのコールバック"""

    def __init__(self):
        # オンメモリでモデルを記録しておく辞書
        self.cv_boosters = {}

    @staticmethod
    def params_to_hash(params):
        """パラメータを元に辞書のキーとなるハッシュを計算する"""
        params_hash = hash(frozenset(params.items()))
        return params_hash

    def get_trained_model(self, params):
        """パラメータをキーとして学習済みモデルを取り出す"""
        params_hash = self.params_to_hash(params)
        return self.cv_boosters[params_hash]

    def __call__(self, env):
        """LightGBM の各ラウンドで呼ばれるコールバック"""
        # 学習に使うパラメータをハッシュ値に変換する
        params_hash = self.params_to_hash(env.params)
        # 初登場のパラメータならモデルへの参照を保持する
        if params_hash not in self.cv_boosters:
            self.cv_boosters[params_hash] = env.model


def main():
    dataset = datasets.load_breast_cancer()
    X, y = dataset.data, dataset.target

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

    lgb_train = lgb.Dataset(X_train, y_train)

    folds = StratifiedKFold(n_splits=3, shuffle=True, random_state=42)

    # 学習済みモデルへの参照を保持するためのコールバック
    checkpoint_cb = TunerCVCheckpointCallback()
    callbacks = [
        checkpoint_cb,
    ]

    lgbm_params = {
        'objective': 'binary',
        'metric': 'binary_logloss',
        'verbosity': -1,
    }
    tuner_cv = lgb.LightGBMTunerCV(
        lgbm_params, lgb_train,
        num_boost_round=1000,
        early_stopping_rounds=100,
        verbose_eval=20,
        folds=folds,
        callbacks=callbacks,
    )

    tuner_cv.run()

    # NOTE: 念のためハッシュの衝突に備えて Trial の数と学習済みモデルの数を比較しておく
    assert len(checkpoint_cb.cv_boosters) == len(tuner_cv.study.trials) - 1

    # 最も良かったパラメータをキーにして学習済みモデルを取り出す
    cv_booster = checkpoint_cb.get_trained_model(tuner_cv.best_params)
    # Averaging でホールドアウト検証データを予測する
    y_pred_proba_list = cv_booster.predict(X_test,
                                           num_iteration=cv_booster.best_iteration)
    y_pred_proba_avg = np.array(y_pred_proba_list).mean(axis=0)
    y_pred = np.where(y_pred_proba_avg > 0.5, 1, 0)
    accuracy = accuracy_score(y_test, y_pred)
    print('Averaging accuracy:', accuracy)


if __name__ == '__main__':
    main()

ちなみに、学習に使ったパラメータをハッシュにしてモデルへの参照のキーにしている。そのため、非常に低い確率だけどハッシュが衝突する可能性を考えて assert 文を入れた。

上記をファイルに保存して実行してみよう。動作確認のために、ホールドアウトしておいたデータを、取り出したモデル (実体は _CVBooster) を使って推論させている。

$ python tunercv2.py 
/Users/amedama/.virtualenvs/py37/lib/python3.7/site-packages/optuna/_experimental.py:83: ExperimentalWarning: LightGBMTunerCV is experimental (supported from v1.5.0). The interface can change in the future.
  ExperimentalWarning,
feature_fraction, val_score: inf:   0%|                   | 0/7 [00:00<?, ?it/s][20] cv_agg's binary_logloss: 0.191389 + 0.0261761
[40]  cv_agg's binary_logloss: 0.133308 + 0.0423785
[60]   cv_agg's binary_logloss: 0.119175 + 0.0400074

...

min_data_in_leaf, val_score: 0.092142: 100%|######| 5/5 [00:01<00:00,  2.51it/s]
Averaging accuracy: 0.972027972027972