TensorFlowのMetricsやLossesを使ってXGBoostを学習する

f:id:astamuse:20210119180122p:plain

こんにちは、2020年3月にastamuseに入社したYNと申します。本ブログを書くのははじめてです。どうぞよろしくお願いいたします。

今回は個人的な趣味で、TensorFlowのMetricsやLossesを使ってXGBoostを学習する方法について紹介したいと思います。

動機

GPUの普及もあってか、XGBoostのようなGBDT (Gradient Boosting Decision Tree) 系のライブラリでもGPUサポートがされるようになりました。ところが、XGBoostのV1.3.0のドキュメントを見ると例えばaucprはGPU対応されていません。また、多くのLossがXGBoostには実装されていますが、近年提案されたLossが実装されているとは限りません。

NN (Neural Network) の流行りもあり、提案されるLossはTensorFlowのようなNN系のライブラリで実装されたものが公開されることが多いです。したがって、TensorFlowで実装されたMetricsやLossesを流用できると簡単にGPU対応できる嬉しさがあります。また、Lossの場合はTensorFlowの自動微分の機能を使うことで、わざわざ手計算でgradientやhessianを計算する必要がなくなるので、私のように数学に自信がなくても手軽に実験できるようになります。私が知る限りでは、過去のKaggleのcompetitionでも少し特殊なLossに対して、TensorFlowを使ってgradientの計算を自動で行ったNotebookもあります。

今回紹介する実験では、Google Colabを使っています。実験環境は以下の通りです。

ハードウェア

CPU: Intel(R) Xeon(R) CPU @ 2.20GHz, GPU: Tesla T4

ライブラリ

numpy==1.19.5
pandas==1.1.5
scikit-learn==0.22.2.post1
tensorflow==2.4.0
tensorflow-addons==0.8.3
xgboost==1.3.0.post0

実験

今回は、ダミーデータで2値分類タスクを解くことを考えます。

はじめに、ベースラインとしてXGBoostに元から実装されているBCE (Binary CrossEntropy) をLoss, PR-AUCをMetricとしたモデルを学習します。

次に、XGBoost 1.3.0ではGPU対応されていないPR-AUCをTensorFlowを使って実装します。

最後に、TensorFlow-addonsのFocal Lossを使ってモデルを学習します。

学習データ・検証データの作成

import numpy as np
import pandas as pd
from sklearn.datasets import make_classification
from sklearn.metrics import average_precision_score
from sklearn.model_selection import train_test_split
import xgboost as xgb


X, y = make_classification(
    n_samples=10_000_000, 
    n_features=20, 
    n_classes=2, 
    random_state=42,
)
X_train, X_valid, y_train, y_valid = train_test_split(X, y, test_size=0.5, random_state=42)

dtrain = xgb.DMatrix(X_train, label=y_train)
dvalid = xgb.DMatrix(X_valid, label=y_valid)

1. ベースライン

まずは、XGBoostの公式の実装を使って、モデルを学習します。学習は以下のようなコードで行いました。

params = {
    'tree_method': 'gpu_hist', 
    'objective': 'binary:logistic',
    'eval_metric': 'aucpr',
    'seed': 42,
}

model = xgb.train(
    params,
    dtrain=dtrain,
    num_boost_round=10,
    evals=[(dvalid, 'valid')],
)
print('PR-AUC: ', average_precision_score(y_valid, model.predict(dvalid)))

PR-AUCは0.9832でした。また、マジックコマンドの%%timeで実行時間を計測したところ、Wall timeは7.47sでした。

2. Metricsの実装

以下のようなclassを作成して、TensorFlowのMetricsクラスを使ってXGBoostのMetricの計算を行えるようにしました。

from typing import Callable, Tuple

import tensorflow as tf


class XGBCustomMetricForTF():
    def __init__(self, tf_metric: tf.keras.metrics.Metric, preprocessor: Callable=None, name: str=None) -> None:
        self.tf_metric = tf_metric
        self.preprocessor = preprocessor
        if name is not None:
            self.name = name
        elif hasattr(tf_metric, '__name__'):
            self.name = tf_metric.__name__
        else:
            self.name = 'custom_metric'

    def __call__(self, y_pred: np.ndarray, dtrain: xgb.DMatrix) -> Tuple[str, np.ndarray]:
        y_true = dtrain.get_label()
        if self.preprocessor is not None:
            y_pred = self.preprocessor(y_pred)
        self.tf_metric.reset_states()
        self.tf_metric.update_state(y_true, y_pred)
        return self.name, self.tf_metric.result().numpy()

今回、PR-AUCの計算には、tf.keras.metrics.AUCを使います。このクラスは厳密な値ではなく、近似値になっていることが注意点です。num_thresholdsを大きくすると近似精度が向上しますが、計算時間が増加します。Ealry Stoppingする位置の決定など、厳密な精度よりも速度が要求される場合には有用かもしれません。

params = {
    'tree_method': 'gpu_hist', 
    'objective': 'binary:logistic',
    'disable_default_eval_metric': 1,
    'seed': 42,
}

custom_metric = XGBCustomMetricForTF(
    tf.keras.metrics.AUC(curve='PR', num_thresholds=100),
    preprocessor=lambda x: 1.0 / (1.0 + np.exp(-x)),
)
model = xgb.train(
    params,
    dtrain=dtrain,
    num_boost_round=10,
    feval=custom_metric,
    evals=[(dvalid, 'valid')],
)
print('PR-AUC: ', average_precision_score(y_valid, model.predict(dvalid)))

PR-AUCは0.9832, Wall timeは5.62sでした。ベースラインでは7.47sだったので、約25%短くなっていますね。 (おそらく、実情はGPUで計算することより、AUCを近似計算することによる高速化のほうが大きいでしょう。)

3. Lossesの実装

Lossesの実装には、以下のようなclassを作成しました。

from typing import Callable, Tuple

import tensorflow as tf


class XGBCustomSingleOutputLossForTF():
    def __init__(self, tf_loss: tf.keras.losses.Loss, preprocessor: Callable=None, name: str=None) -> None:
        self.tf_loss = tf_loss
        self.preprocessor = preprocessor
        if name is not None:
            self.name = name
        elif hasattr(tf_loss, '__name__'):
            self.name = tf_loss.__name__
        else:
            self.name = 'custom_metric'

    def __call__(self, y_pred: np.ndarray, dtrain: xgb.DMatrix) -> Tuple[np.ndarray, np.ndarray]:
        y_true = dtrain.get_label().reshape(-1, 1) # single output
        if self.preprocessor is not None:
            y_pred = self.preprocessor(y_pred)
        y_pred = y_pred.reshape(-1, 1) # single output
        y_pred = tf.convert_to_tensor(y_pred)
        
        with tf.GradientTape() as t1:
            t1.watch(y_pred)
            with tf.GradientTape() as t2:
                t2.watch(y_pred)
                loss = self.tf_loss(y_true, y_pred)
            grad = t2.gradient(loss, y_pred)
        hess = t1.gradient(grad, y_pred)

        return grad.numpy(), hess.numpy()

このclassを使ってモデルの学習を行います。TensorFlowのリポジトリにはFocal Lossがないので、TensorFlow addonsを使います。

import tensorflow_addons as tfa


params = {
    'tree_method': 'gpu_hist', 
    'objective': 'binary:logistic',
    'eval_metric': 'aucpr',
    'seed': 42,
}

custom_obj = XGBCustomSingleOutputLossForTF(
    tfa.losses.SigmoidFocalCrossEntropy(reduction=tf.keras.losses.Reduction.NONE),
    preprocessor=lambda x: 1.0 / (1.0 + np.exp(-x)),
)
model = xgb.train(
    params,
    dtrain=dtrain,
    num_boost_round=10,
    obj=custom_obj,
    evals=[(dvalid, 'valid')],
)
print('PR-AUC: ', average_precision_score(y_valid, model.predict(dvalid)))