Python: XGBoost の cv() 関数から学習済みモデルを取り出す

今回は、以下のエントリを XGBoost で焼き直したもの。つまり、XGBoost でも cv() 関数から学習済みモデルを取り出して Fold Averaging してみようという話。

blog.amedama.jp

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

$ sw_vers                   
ProductName:    Mac OS X
ProductVersion: 10.14.6
BuildVersion:   18G5033
$ python -V              
Python 3.7.7
$ pip list | grep xgboost               
xgboost                    1.1.1

　下準備

必要なパッケージをインストールしておく。

$ pip install xgboost scikit-learn numpy

学習済みモデルを取り出して Fold Averaging してみる

早速、以下にサンプルコードを示す。乳がんデータセットをホールドアウトして、一方のデータで学習して、他方のデータを Fold Averaging している。

実装方法としては、LightGBM と同じようにコールバックを使って学習済みモデルへの参照を引っ張り出した。 cv() 関数のコールバックには cvfolds というパラメータ名で xgboost.training.CVPack のリストが渡される。あとは CVPack#bst という名前で Booster オブジェクトにアクセスするだけ。

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

import inspect

import numpy as np
import xgboost as xgb
from sklearn import datasets
from sklearn.metrics import accuracy_score
from sklearn.model_selection import StratifiedKFold
from sklearn.model_selection import train_test_split


class ModelExtractionCallback:
    """XGBoost の cv() 関数から学習済みモデルを取り出すためのコールバック"""

    def __init__(self):
        self.cvfolds = None

    def __call__(self, env):
        # コールバックの呼び出しで xgboost.training.CVPack のリストが得られる
        if self.cvfolds is None:
            self.cvfolds = env.cvfolds


class BoostersProxy:
    """コールバックから得られた CVPack のリストを使って Fold Averaging をやりやすくするクラス

    実用上は ModelExtractionCallback とニコイチしちゃっても良いけど一応分けた"""

    def __init__(self, cvfolds):
        self._cvfolds = cvfolds

    def __getattr__(self, name):
        def _wrap(*args, **kwargs):
            ret = []
            for cvpack in self._cvfolds:
                # それぞれの Booster から指定されたアトリビュートを取り出す
                attr = getattr(cvpack.bst, name)
                if inspect.ismethod(attr):
                    # オブジェクトがメソッドなら呼び出した結果を入れる
                    res = attr(*args, **kwargs)
                    ret.append(res)
                else:
                    # それ以外ならそのまま入れる
                    ret.append(attr)
            return ret
        return _wrap


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,
                                                        test_size=0.33,
                                                        shuffle=True,
                                                        random_state=42,
                                                        stratify=y)

    # XGBoost のデータセット表現に直す
    dtrain = xgb.DMatrix(X_train, label=y_train)
    dtest = xgb.DMatrix(X_test, label=y_test)

    # データの分割に使うコンテキスト
    folds = StratifiedKFold(n_splits=5,
                            shuffle=True,
                            random_state=42)
    # 学習に使うパラメータ
    xgb_params = {
        'objective': 'binary:logistic',
        'eval_metric': 'logloss',
    }

    # モデルを取り出すのに使うコールバック
    extraction_cb = ModelExtractionCallback()
    callbacks = [
        extraction_cb,
    ]

    # 交差検証する
    xgb.cv(xgb_params,
           dtrain,
           num_boost_round=1000,
           early_stopping_rounds=100,
           folds=folds,
           verbose_eval=True,
           # コールバックを渡す
           callbacks=callbacks,
           )

    # コールバックから学習済みモデルを取り出してプロキシにくべる
    proxy = BoostersProxy(cvfolds=extraction_cb.cvfolds)

    # ホールドアウトしておいた検証データを Fold Averaging で予測する
    y_pred_proba_list = proxy.predict(dtest)
    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('Fold averaging accuracy:', accuracy)


if __name__ == '__main__':
    main()

上記を実行してみよう。

$ python xgbcv.py 
[0]    train-logloss:0.46506+0.00234  test-logloss:0.48322+0.01040
[1]    train-logloss:0.33602+0.00369  test-logloss:0.36709+0.01240
[2]    train-logloss:0.25204+0.00529  test-logloss:0.29438+0.01717

...

[130]  train-logloss:0.00727+0.00008  test-logloss:0.11461+0.07024
[131]  train-logloss:0.00727+0.00008  test-logloss:0.11468+0.07019
[132]  train-logloss:0.00727+0.00008  test-logloss:0.11473+0.07017
Fold averaging accuracy: 0.9627659574468085