Python: CatBoost を GPU で学習させる

勾配ブースティング決定木を扱うフレームワークの CatBoost は、GPU を使った学習ができる。 GPU を使うと、CatBoost の特徴的な決定木の作り方 (Symmetric Tree) も相まって、学習速度の向上が見込める場合があるようだ。今回は、それを試してみる。

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

$ cat /etc/*-release
DISTRIB_ID=Ubuntu
DISTRIB_RELEASE=20.04
DISTRIB_CODENAME=focal
DISTRIB_DESCRIPTION="Ubuntu 20.04.1 LTS"
NAME="Ubuntu"
VERSION="20.04.1 LTS (Focal Fossa)"
ID=ubuntu
ID_LIKE=debian
PRETTY_NAME="Ubuntu 20.04.1 LTS"
VERSION_ID="20.04"
HOME_URL="https://www.ubuntu.com/"
SUPPORT_URL="https://help.ubuntu.com/"
BUG_REPORT_URL="https://bugs.launchpad.net/ubuntu/"
PRIVACY_POLICY_URL="https://www.ubuntu.com/legal/terms-and-policies/privacy-policy"
VERSION_CODENAME=focal
UBUNTU_CODENAME=focal
$ python3 -V
Python 3.8.2
$ python3 -m pip list | grep -i catboost
catboost               0.24

使用するハードウェアは Google Compute Engine の N1 Standard 2 インスタンスに NVIDIA Tesla T4 を 1 台アタッチしている。

$ grep "model name" /proc/cpuinfo | head -n 1
model name   : Intel(R) Xeon(R) CPU @ 2.30GHz
$ grep processor /proc/cpuinfo | wc -l
2
$ nvidia-smi
Sat Aug 22 07:05:51 2020       
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 450.51.06    Driver Version: 450.51.06    CUDA Version: 11.0     |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|                               |                      |               MIG M. |
|===============================+======================+======================|
|   0  Tesla T4            On   | 00000000:00:04.0 Off |                    0 |
| N/A   47C    P8    10W /  70W |     70MiB / 15109MiB |      0%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+
                                                                               
+-----------------------------------------------------------------------------+
| Processes:                                                                  |
|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |
|        ID   ID                                                   Usage      |
|=============================================================================|
|    0   N/A  N/A       889      G   /usr/lib/xorg/Xorg                 59MiB |
|    0   N/A  N/A       983      G   /usr/bin/gnome-shell               10MiB |
+-----------------------------------------------------------------------------+

下準備

はじめに、CatBoost とその他使うパッケージをインストールしておく。

$ python3 -m pip install catboost scikit-learn

なお、CatBoost では GPU のリソースを CUDA で扱うので、使用するマシンにはあらかじめ CUDA をインストールしておく。今回は CUDA をインストールする部分については手順から省略する。次のスニペットを実行して、結果が 0 でなければ GPU のリソースが CatBoost から見えていることがわかる。

$ python3 -c "from catboost.utils import get_gpu_device_count; print(get_gpu_device_count())"
1

CatBoost を GPU を使って学習する

以下のサンプルコードでは、擬似的に作った二値分類のデータセットを CatBoost で学習させている。ポイントは、学習するときに渡す辞書のパラメータに task_type というキーで GPU を指定するところ。 CatBoost から GPU のリソースが認識できていれば、これだけで GPU を使った学習ができる。

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

import sys
import time
import logging
from contextlib import contextmanager

from catboost import CatBoost
from catboost import Pool
from sklearn.datasets import make_classification
from sklearn.model_selection import train_test_split
from sklearn.metrics import log_loss


LOGGER = logging.getLogger(__name__)


@contextmanager
def timeit():
    """処理にかかった時間を計測してログに出力するコンテキストマネージャ"""
    start = time.time()
    yield
    end = time.time()
    elapsed = end - start
    LOGGER.info(f'Elapsed Time: {elapsed:.2f} sec')


def main():
    logging.basicConfig(level=logging.INFO,
                        stream=sys.stderr,
                        )

    # 疑似的な教師信号を作るためのパラメータ
    dist_args = {
        # データ点数
        'n_samples': 100_000,
        # 次元数
        'n_features': 1_000,
        # その中で意味のあるもの
        'n_informative': 100,
        # 重複や繰り返しはなし
        'n_redundant': 0,
        'n_repeated': 0,
        # タスクの難易度
        'class_sep': 0.65,
        # 二値分類問題
        'n_classes': 2,
        # 生成に用いる乱数
        'random_state': 42,
        # 特徴の順序をシャッフルしない (先頭の次元が informative になる)
        'shuffle': False,
    }
    # 教師データを作る
    train_x, train_y = make_classification(**dist_args)
    # データセットを学習用と検証用に分割する
    x_tr, x_val, y_tr, y_val = train_test_split(train_x, train_y,
                                                test_size=0.3,
                                                shuffle=True,
                                                random_state=42,
                                                stratify=train_y)
    # CatBoost が扱うデータセットの形式に直す
    train_pool = Pool(x_tr, label=y_tr)
    valid_pool = Pool(x_val, label=y_val)
    # 学習用のパラメータ
    params = {
        # タスク設定と損失関数
        'loss_function': 'Logloss',
        # 学習率
        'learning_rate': 0.02,
        # 学習ラウンド数
        'num_boost_round': 5_000,
        # 検証用データの損失が既定ラウンド数減らなかったら学習を打ち切る
        # NOTE: ラウンド数を揃えたいので今回は使わない
        # 'early_stopping_rounds': 100,
        # 乱数シード
        'random_state': 42,
        # 学習に GPU を使う場合
        'task_type': 'GPU',
    }
    # モデルを学習する
    model = CatBoost(params)
    with timeit():
        model.fit(train_pool,
                  eval_set=valid_pool,
                  verbose_eval=100,
                  use_best_model=True,
                  )
    # 検証用データを分類する
    y_pred = model.predict(valid_pool,
                           prediction_type='Probability')
    # ロジスティック損失を確認する
    metric = log_loss(y_val, y_pred)
    LOGGER.info(f'Validation Metric: {metric}')


if __name__ == '__main__':
    main()

上記を実行してみよう。

$ python3 catgpubench.py
0: learn: 0.6917674   test: 0.6918382  best: 0.6918382 (0)   total: 41.1ms  remaining: 3m 25s
100:   learn: 0.5880012   test: 0.5928815  best: 0.5928815 (100) total: 2.34s   remaining: 1m 53s
200:   learn: 0.5159286   test: 0.5238258  best: 0.5238258 (200) total: 4.52s   remaining: 1m 47s

...

4800:  learn: 0.0799767   test: 0.1156692  best: 0.1156685 (4799)    total: 1m 36s   remaining: 4.01s
4900:  learn: 0.0790359   test: 0.1150905  best: 0.1150903 (4899)    total: 1m 38s   remaining: 1.99s
4999:  learn: 0.0780408   test: 0.1144641  best: 0.1144641 (4999)    total: 1m 40s   remaining: 0us
bestTest = 0.1144640706
bestIteration = 4999
INFO:__main__:Elapsed Time: 109.65 sec
INFO:__main__:Validation Metric: 0.11446408301027777

ちゃんと GPU を使って学習できた。 GPU が使われているかは nvidia-smi などでリソースの状態を確認すると良いと思う。

CPU を使った学習と比べてみる

一応、CPU を使って学習したときと比べてみよう。ただ、先ほど使ったインスタンスは CPU のコア数があまりにも少ない。そのため、似たような値段で借りられる N1 Standard 16 インスタンスを使って比較する。

ハードウェアの環境的には次のとおり。

$ grep "model name" /proc/cpuinfo | head -n 1
model name   : Intel(R) Xeon(R) CPU @ 2.30GHz
$ grep processor /proc/cpuinfo | wc -l
16

あらかじめ前置きしておくと、GPU を使って高速化が見込めるかは、データセットの特性や学習時のオプション、ハードウェアなど様々なパラメータに依存する。なので、今回の内容はあくまで「特定の環境で試したときにこうなった」という結果に過ぎない。

ソースコードを編集して、学習時のパラメータで GPU を使わないようにする。

    # 学習用のパラメータ
    params = {
        # タスク設定と損失関数
        'loss_function': 'Logloss',
        # 学習率
        'learning_rate': 0.02,
        # 学習ラウンド数
        'num_boost_round': 5_000,
        # 検証用データの損失が既定ラウンド数減らなかったら学習を打ち切る
        # NOTE: ラウンド数を揃えたいので今回は使わない
        # 'early_stopping_rounds': 100,
        # 乱数シード
        'random_state': 42,
        # 学習に GPU を使う場合
        # 'task_type': 'GPU',
    }

そして、実行しよう。

$ python3 catcpubench.py
0: learn: 0.6916098   test: 0.6916659  best: 0.6916659 (0)   total: 239ms    remaining: 19m 56s
100:   learn: 0.5917145   test: 0.5961182  best: 0.5961182 (100) total: 7.93s   remaining: 6m 24s
200:   learn: 0.5218843   test: 0.5286355  best: 0.5286355 (200) total: 15.5s   remaining: 6m 10s

...

4800:  learn: 0.0643858   test: 0.1035075  best: 0.1035075 (4800)    total: 5m 40s   remaining: 14.1s
4900:  learn: 0.0629871   test: 0.1023799  best: 0.1023799 (4900)    total: 5m 47s   remaining: 7.01s
4999:  learn: 0.0618029   test: 0.1015037  best: 0.1015037 (4999)    total: 5m 53s   remaining: 0us

bestTest = 0.1015037231
bestIteration = 4999

INFO:__main__:Elapsed Time: 356.12 sec
INFO:__main__:Validation Metric: 0.10150372305811575