CUBE SUGAR CONTAINER

技術系のこと書きます。

Python: LightGBM v4.0 の CUDA 実装を試す

LightGBM のバージョン 4.0.0 が 2023-07-14 にリリースされた。 このリリースは久しぶりのメジャーアップデートで、様々な改良が含まれている。 詳細については、以下のリリースノートで確認できる。

github.com

リリースの大きな目玉として CUDA を使った学習の実装が全面的に書き直されたことが挙げられる。 以前の LightGBM は、GPU を学習に使う場合でも、その計算リソースを利用できる範囲が限られていた。 それが、今回の全面的な刷新によって、利用の範囲が拡大されたとのこと。

ただし、PyPI で配布されている Linux 向け Wheel ファイルは CUDA での学習に対応していない。 対応しているのは CPU と、GPU でも OpenCL の API を使ったもの。 そのため、もし CUDA を使った学習を利用したい場合には自分で Wheel をビルドする必要がある。

使った環境は次のとおり。 OS が Ubuntu 22.04 LTS で、CPU が Intel Core i7-12700、GPU が NVIDIA GeForce RTX 3060 のマシンになる。

$ cat /etc/lsb-release 
DISTRIB_ID=Ubuntu
DISTRIB_RELEASE=22.04
DISTRIB_CODENAME=jammy
DISTRIB_DESCRIPTION="Ubuntu 22.04.2 LTS"
$ uname -srm
Linux 5.15.0-76-generic x86_64
$ pip --version
pip 23.1.2 from /usr/local/lib/python3.10/dist-packages/pip (python 3.10)
$ cat /proc/cpuinfo  | grep "model name" | head -n 1
model name  : 12th Gen Intel(R) Core(TM) i7-12700
$ nvidia-smi
Fri Jul 14 18:21:47 2023       
+---------------------------------------------------------------------------------------+
| NVIDIA-SMI 535.54.03              Driver Version: 535.54.03    CUDA Version: 12.2     |
|-----------------------------------------+----------------------+----------------------+
| 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  NVIDIA GeForce RTX 3060        On  | 00000000:01:00.0 Off |                  N/A |
|  0%   51C    P8              14W / 170W |     19MiB / 12288MiB |      0%      Default |
|                                         |                      |                  N/A |
+-----------------------------------------+----------------------+----------------------+
                                                                                         
+---------------------------------------------------------------------------------------+
| Processes:                                                                            |
|  GPU   GI   CI        PID   Type   Process name                            GPU Memory |
|        ID   ID                                                             Usage      |
|=======================================================================================|
|    0   N/A  N/A      1095      G   /usr/lib/xorg/Xorg                            9MiB |
|    0   N/A  N/A      1240      G   /usr/bin/gnome-shell                          2MiB |
+---------------------------------------------------------------------------------------+

もくじ

下準備

まずは Wheel のビルドに必要なパッケージなどをインストールする。

$ sudo apt-get install \
    --no-install-recommends \
    git \
    cmake \
    build-essential \
    libboost-dev \
    libboost-system-dev \
    libboost-filesystem-dev \
    python3 \
    wget \
    unzip

次に PIP をインストールする。 APT のパッケージを使わないのは、古いバージョンではサポートされていないオプションを利用するため。

$ wget -O - https://bootstrap.pypa.io/get-pip.py | sudo python3

そして本題となる CUDA を有効にした Wheel をビルドしてインストールする手順が以下になる。 まず --no-binary で PyPI の Wheel をインストールせず、ソースコード配布物を自身でビルドする。 また、過去のキャッシュが効いてしまわないように --no-cache も指定しておく。 そして、CUDA を有効にするために --config-settingscmake.define.USE_CUDA=ON を指定する。

$ pip install lightgbm \
    --no-binary lightgbm \
    --no-cache lightgbm \
    --config-settings=cmake.define.USE_CUDA=ON

しばらくすれば Wheel のビルドとインストールが完了するはず。

あとは、今回他に利用するパッケージをインストールしておく。

$ pip install scikit-learn polars

動作を確認するのに使うデータセットとして HIGGS をダウンロードしておく。 このデータセットは二値分類のタスクで、サイズが大きくてあらかじめすべてのカラムが数値なので前処理が必要ないという特徴がある。

$ wget https://archive.ics.uci.edu/static/public/280/higgs.zip
$ unzip higgs.zip
$ gunzip HIGGS.csv.gz
$ du -m HIGGS.csv 
7664   HIGGS.csv

サンプルコード

以下にサンプルコードを示す。 基本的には CPU を学習に使うコードと同じものが利用できるけど、パラメータによっては CPU でしか使えないものがある 1。 今回利用しているパラメータはほとんどがデフォルトなので特に影響がない。 学習に CUDA を利用する場合には、LightGBM の学習パラメータで "device" のデフォルトが "cpu" であるところを "cuda" にすれば良い。

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

import logging
import contextlib
import time

import polars as pl
import lightgbm as lgb
from sklearn.model_selection import train_test_split
from sklearn.metrics import roc_auc_score

LOG = logging.getLogger(__name__)


@contextlib.contextmanager
def timer():
    """実行時間の計測に使うコンテキストマネージャ"""
    start_time = time.time()
    yield
    end_time = time.time()
    elapsed_time = end_time - start_time
    LOG.info("elapsed time: %.3f sec", elapsed_time)


def main():
    # INFO レベル以上のログを出力する
    logging.basicConfig(
        level=logging.INFO,
    )

    # HIGGS データセットを読み込む
    columns = {
        "class_label": pl.Float32,
        "jet_1_b-tag": pl.Float64,
        "jet_1_eta": pl.Float64,
        "jet_1_phi": pl.Float64,
        "jet_1_pt": pl.Float64,
        "jet_2_b-tag": pl.Float64,
        "jet_2_eta": pl.Float64,
        "jet_2_phi": pl.Float64,
        "jet_2_pt": pl.Float64,
        "jet_3_b-tag": pl.Float64,
        "jet_3_eta": pl.Float64,
        "jet_3_phi": pl.Float64,
        "jet_3_pt": pl.Float64,
        "jet_4_b-tag": pl.Float64,
        "jet_4_eta": pl.Float64,
        "jet_4_phi": pl.Float64,
        "jet_4_pt": pl.Float64,
        "lepton_eta": pl.Float64,
        "lepton_pT": pl.Float64,
        "lepton_phi": pl.Float64,
        "m_bb": pl.Float64,
        "m_jj": pl.Float64,
        "m_jjj": pl.Float64,
        "m_jlv": pl.Float64,
        "m_lv": pl.Float64,
        "m_wbb": pl.Float64,
        "m_wwbb": pl.Float64,
        "missing_energy_magnitude": pl.Float64,
        "missing_energy_phi": pl.Float64,
    }
    df = pl.read_csv(
        source="HIGGS.csv",
        columns=list(columns.keys()),
        dtypes=columns,
    )

    # 説明変数と目的変数に分けて NumPy 配列に直す
    x = df.select(pl.exclude("class_label")).to_numpy()
    y = df.get_column("class_label").cast(pl.Int8).to_numpy()

    # データセットを学習用とテスト用に分けておく
    train_x, test_x, train_y, test_y = train_test_split(
        x, y,
        test_size=0.3,
        stratify=y,
        shuffle=True,
        random_state=42,
    )

    # 学習用データセットを、さらに学習用と評価用に分割する
    train_x, valid_x, train_y, valid_y = train_test_split(
        train_x, train_y,
        test_size=0.2,
        stratify=train_y,
        shuffle=True,
        random_state=42,
    )

    # LightGBM のデータセットオブジェクトにする
    lgb_train = lgb.Dataset(train_x, train_y)
    lgb_valid = lgb.Dataset(valid_x, valid_y, reference=lgb_train)

    # 学習用のパラメータ
    lgb_params = {
        # 目的関数
        "objective": "binary",
        # 評価指標
        "metric": "auc",
        # 乱数シード
        "seed": 42,
        # 学習に使うデバイス: cpu / cuda
        "device": "cpu",
    }
    # コールバック関数
    lgb_callbacks = [
        # 100 イテレーションごとにメトリックを出力する
        lgb.log_evaluation(
            period=100,
        ),
    ]

    # 1,000 イテレーションの実行に要する時間を計測する
    with timer():
        booster = lgb.train(
            params=lgb_params,
            train_set=lgb_train,
            num_boost_round=1_000,
            valid_sets=[lgb_valid],
            callbacks=lgb_callbacks,
        )

    # テスト用のデータセットで評価指標を確認する
    y_pred = booster.predict(test_x)
    test_metric = roc_auc_score(test_y, y_pred)
    LOG.info("ROC-AUC score: %.6f", test_metric)


if __name__ == "__main__":
    main()

CPU を使った学習を試す

まずは CPU で学習する場合を実行してみよう。 先ほどのサンプルコードをそのまま実行するだけ。

$ python3 example.py 
[LightGBM] [Info] Number of positive: 3264308, number of negative: 2895691
[LightGBM] [Warning] Auto-choosing col-wise multi-threading, the overhead of testing was 0.252465 seconds.
You can set `force_col_wise=true` to remove the overhead.
[LightGBM] [Info] Total Bins 6132
[LightGBM] [Info] Number of data points in the train set: 6159999, number of used features: 28
[LightGBM] [Info] [binary:BoostFromScore]: pavg=0.529920 -> initscore=0.119824
[LightGBM] [Info] Start training from score 0.119824
[100]  valid_0's auc: 0.811524
[200] valid_0's auc: 0.819294
[300]  valid_0's auc: 0.822628
[400] valid_0's auc: 0.825188
[500]  valid_0's auc: 0.827045
[600] valid_0's auc: 0.828482
[700]  valid_0's auc: 0.829784
[800] valid_0's auc: 0.831063
[900]  valid_0's auc: 0.832142
[1000]    valid_0's auc: 0.833035
INFO:__main__:elapsed time: 414.358 sec
INFO:__main__:ROC-AUC score: 0.833369

今回の環境では、学習に約 414 秒かかった。 ホールド・アウトしておいたデータに対する性能は ROC-AUC で 0.833369 となっている。

CUDA を使った学習を試す

続いて CUDA を使って学習してみよう。 先ほどのサンプルコードの "device": "cpu""device": "cuda" に書き換えて実行する。 ROC-AUC の計算が CUDA には対応していないので CPU の処理にフォールバックしているけど、メトリックの計算に過ぎないので特に問題はないはず。

$ python3 benchmark.py
[LightGBM] [Warning] Using sparse features with CUDA is currently not supported.
[LightGBM] [Info] Number of positive: 3264308, number of negative: 2895691
[LightGBM] [Warning] Metric auc is not implemented in cuda version. Fall back to evaluation on CPU.
[LightGBM] [Info] Total Bins 6132
[LightGBM] [Info] Number of data points in the train set: 6159999, number of used features: 28
[LightGBM] [Warning] Metric auc is not implemented in cuda version. Fall back to evaluation on CPU.
[LightGBM] [Info] [binary:BoostFromScore]: pavg=0.529920 -> initscore=0.119824
[LightGBM] [Info] Start training from score 0.119824
[100]  valid_0's auc: 0.811524
[200] valid_0's auc: 0.819294
[300]  valid_0's auc: 0.822628
[400] valid_0's auc: 0.825175
[500]  valid_0's auc: 0.826881
[600] valid_0's auc: 0.82841
[700]  valid_0's auc: 0.829691
[800] valid_0's auc: 0.83085
[900]  valid_0's auc: 0.832043
[1000]    valid_0's auc: 0.833091
INFO:__main__:elapsed time: 79.020 sec
INFO:__main__:ROC-AUC score: 0.833390

今回の環境では、学習に約 79 秒かかった。 ホールド・アウトしておいたデータに対する性能は ROC-AUC で 0.833390 となっている。

CPU を学習に使う場合に比べて、要する時間が約 1/5 に短縮されている。 そして、ホールド・アウトしたデータに対する性能も悪化しているようには見えない。

まとめ

今回は LightGBM v4.0 で書き直された CUDA の実装を使った学習を試してみた。 あくまで今回の環境においてはであるが、学習にかかる時間が大幅に短縮できること、そして汎化性能が悪化しないように見えることが確認できた。