CUBE SUGAR CONTAINER

技術系のこと書きます。

Python: LightGBM の学習率を動的に制御する

LightGBM の学習率は基本的に低い方が最終的に得られるモデルの汎化性能が高くなることが経験則として知られている。 しかしながら、学習率が低いとモデルの学習に多くのラウンド数、つまり計算量を必要とする。 そこで、今回は学習率を学習の過程において動的に制御するコールバックを実装してみた。

きっかけは以下のツイートを見たこと。

なるほど面白そう。

下準備

使用するライブラリをあらかじめインストールしておく。

$ pip install lightgbm seaborn scikit-learn

学習率を動的に制御するコールバック

早速だけど、以下が学習率を動的に制御するコールバックを実装したサンプルコードとなる。 コールバックの本体は LrSchedulingCallback というクラスで実装している。 このクラスをインスタンス化するときに、制御方法を記述した関数を渡す。 以下であれば sample_scheduler_func() という名前で定義した。 この関数は学習の履歴などを元に新たな学習率を決めて返すインターフェースとなっている。 今回はお試しとして 10 ラウンドごとに学習率を下限の 0.01 まで半減させ続けるという単純な戦略を記述してみた。 もちろん、これがベストというわけではなくて、あくまでサンプルとして簡単なものを書いてみたに過ぎない。 なお、EarlyStopping していないのは学習の過程を最後まで観察するため。

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

import numpy as np
import lightgbm as lgb
import seaborn as sns
from matplotlib import pyplot as plt
from sklearn.model_selection import StratifiedKFold


def sample_scheduler_func(current_lr, eval_history, best_round, is_higher_better):
    """次のラウンドで用いる学習率を決定するための関数 (この中身を好きに改造する)

    :param current_lr: 現在の学習率 (指定されていない場合の初期値は None)
    :param eval_history: 検証用データに対する評価指標の履歴
    :param best_round: 現状で最も評価指標の良かったラウンド数
    :param is_higher_better: 高い方が性能指標として優れているか否か
    :return: 次のラウンドで用いる学習率

    NOTE: 学習を打ち切りたいときには callback.EarlyStopException を上げる
    """
    # 学習率が設定されていない場合のデフォルト
    current_lr = current_lr or 0.05

    # 試しに 20 ラウンド毎に学習率を半分にしてみる
    if len(eval_history) % 20 == 0:
        current_lr /= 2

    # 小さすぎるとほとんど学習が進まないので下限も用意する
    min_threshold = 0.001
    current_lr = max(min_threshold, current_lr)

    return current_lr


class LrSchedulingCallback(object):
    """ラウンドごとの学習率を動的に制御するためのコールバック"""

    def __init__(self, strategy_func):
        # 学習率を決定するための関数
        self.scheduler_func = strategy_func
        # 検証用データに対する評価指標の履歴
        self.eval_metric_history = []

    def __call__(self, env):
        # 現在の学習率を取得する
        current_lr = env.params.get('learning_rate')

        # 検証用データに対する評価結果を取り出す (先頭の評価指標)
        first_eval_result = env.evaluation_result_list[0]
        # スコア
        metric_score = first_eval_result[2]
        # 評価指標は大きい方が優れているか否か
        is_higher_better = first_eval_result[3]

        # 評価指標の履歴を更新する
        self.eval_metric_history.append(metric_score)
        # 現状で最も優れたラウンド数を計算する
        best_round_find_func = np.argmax if is_higher_better else np.argmin
        best_round = best_round_find_func(self.eval_metric_history)

        # 新しい学習率を計算する
        new_lr = self.scheduler_func(current_lr=current_lr,
                                     eval_history=self.eval_metric_history,
                                     best_round=best_round,
                                     is_higher_better=is_higher_better)

        # 次のラウンドで使う学習率を更新する
        update_params = {
            'learning_rate': new_lr,
        }
        env.model.reset_parameter(update_params)
        env.params.update(update_params)

    @property
    def before_iteration(self):
        # コールバックは各イテレーションの後に実行する
        return False


def main():
    # Titanic データセットを読み込む
    dataset = sns.load_dataset('titanic')

    # 重複など不要な特徴量は落とす
    X = dataset.drop(['survived',
                      'class',
                      'who',
                      'embark_town',
                      'alive'], axis=1)
    y = dataset.survived

    # カテゴリカル変数を指定する
    categorical_columns = ['pclass',
                           'sex',
                           'embarked',
                           'adult_male',
                           'deck',
                           'alone']
    X = X.astype({c: 'category'
                  for c in categorical_columns})

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

    # コールバックを用意する
    lr_scheduler_cb = LrSchedulingCallback(strategy_func=sample_scheduler_func)
    callbacks = [
        lr_scheduler_cb,
    ]

    # 二値分類を LogLoss で評価する
    lgb_params = {
        'objective': 'binary',
        'metrics': 'binary_logloss',
        'min_data_in_leaf': 10,
    }
    # 5-Fold CV
    skf = StratifiedKFold(n_splits=5,
                          shuffle=True,
                          random_state=42)

    # 動的に学習率を制御した場合
    cv_results = lgb.cv(lgb_params, lgb_train,
                        num_boost_round=500,
                        verbose_eval=1,
                        folds=skf, seed=42,
                        callbacks=callbacks,
                        )
    dynamic_lr = cv_results['binary_logloss-mean']

    # 学習率を 0.1 に固定した場合
    lgb_params.update({'learning_rate': 0.1})
    cv_results = lgb.cv(lgb_params, lgb_train,
                        num_boost_round=500,
                        verbose_eval=1,
                        folds=skf, seed=42,
                        )
    static_lr_0_1 = cv_results['binary_logloss-mean']

    # 学習率を 0.05 に固定した場合
    lgb_params.update({'learning_rate': 0.05})
    cv_results = lgb.cv(lgb_params, lgb_train,
                        num_boost_round=500,
                        verbose_eval=1,
                        folds=skf, seed=42,
                        )
    static_lr_0_05 = cv_results['binary_logloss-mean']

    # 学習率を 0.01 に固定した場合
    lgb_params.update({'learning_rate': 0.01})
    cv_results = lgb.cv(lgb_params, lgb_train,
                        num_boost_round=500,
                        verbose_eval=1,
                        folds=skf, seed=42,
                        )
    static_lr_0_01 = cv_results['binary_logloss-mean']

    # 最小の損失を比較する
    print('min loss value (lr=dynamic):', min(dynamic_lr))
    print('min loss value (lr=0.1):', min(static_lr_0_1))
    print('min loss value (lr=0.05):', min(static_lr_0_05))
    print('min loss value (lr=0.01):', min(static_lr_0_01))

    # 最小の損失が得られたラウンド数を比較する
    print('min loss round (lr=dynamic):', np.argmin(dynamic_lr))
    print('min loss round (lr=0.1):', np.argmin(static_lr_0_1))
    print('min loss round (lr=0.05):', np.argmin(static_lr_0_05))
    print('min loss round (lr=0.01):', np.argmin(static_lr_0_01))

    # グラフにプロットする
    sns.lineplot(np.arange(len(dynamic_lr)),
                 dynamic_lr,
                 label='LR=dynamic')
    sns.lineplot(np.arange(len(static_lr_0_1)),
                 static_lr_0_1,
                 label='LR=0.1')
    sns.lineplot(np.arange(len(static_lr_0_05)),
                 static_lr_0_05,
                 label='LR=0.05')
    sns.lineplot(np.arange(len(static_lr_0_01)),
                 static_lr_0_01,
                 label='LR=0.01')
    plt.title('learning rate control comparison')
    plt.xlabel('rounds')
    plt.ylabel('logloss')
    plt.legend()
    plt.show()


if __name__ == '__main__':
    main()

上記を実行してみる。 最も性能が良かったモデルの損失とラウンド数が表示される。 最適なモデルの性能、学習に要するラウンド数ともに 0.1 固定と 0.01 固定の間にあることが分かる。

$ python dynamiclr.py
...(snip)...
min loss value (lr=dynamic): 0.421097448234137
min loss value (lr=0.1): 0.42265029913071334
min loss value (lr=0.05): 0.4221001657363532
min loss value (lr=0.01): 0.42100303104081405
min loss round (lr=dynamic): 84
min loss round (lr=0.1): 18
min loss round (lr=0.05): 38
min loss round (lr=0.01): 196

そして、各条件における検証用データに対する評価指標の推移をプロットしたグラフが次の通り。 学習率を動的に制御しているパターンは、0.1 固定ほどではないにせよ早く性能が収束していることが分かる。 まあ、とはいえこれくらいなら lr=0.01 ~ 0.05 の間に似たような特性の学習率がいるかもしれない。

いじょう。 こんな上手くいくスケジューラが書けた、みたいな話があったら教えてほしいな。