はじめに
みずほリサーチ＆テクノロジーズ株式会社の@fujineです。
いきなりですがAIエンジニアの皆さん、scikit-learnが実験的にGPUに対応していたこと、ご存知でしょうか？
scikit-learnは機械学習分野における古参パッケージの1つです。多様な機能を提供する一方、FAQにて「GPUに対応する予定はない（キリッ）」と公式宣言しており、scikit-learnが好きな自分としては「勿体無いなぁ」と常々感じていました。
そんな中、何気なくRelease Highlights for 1.2を読んでいたら以下文面を発見！しかも約半年前の2022年12月にリリースされてる…
Experimental Array API support in LinearDiscriminantAnalysis
Experimental support for the Array API specification was added to LinearDiscriminantAnalysis. The estimator can now run on any Array API compliant libraries such as CuPy, a GPU-accelerated array library. For details, see the User Guide.
ということで、そんなに新しくないですが個人的に衝撃的なニュースだったので、本当にGPUに対応しているのかをExample usageに沿って検証しました。
概要
- scikit-learnのver1.2にて、GPUによる学習・推論が実験的に導入された
- 現時点では、まだ1種類の線型判別器（LinearDiscriminantAnalysis）のみに適用可能
- GitHubではissueやPRが活発に提案されており、今後様々な前処理や予測器がGPU上で実行可能になると期待される
検証環境
お手軽にKaggleのNotebook環境を利用します。オプションにて、ACCELERATORはTesla P100、ENVIRONMENTは「Always use latest environment」をそれぞれ選択します。
検証環境のPythonとOSのバージョン、GPUデバイスの情報は以下の通りです。
$ python -V
$ lsb_release -a
$ lspci | grep -i nvidia
Python 3.10.10
No LSB modules are available.
Distributor ID: Ubuntu
Description: Ubuntu 22.04.2 LTS
Release: 22.04
Codename: jammy
00:04.0 3D controller: NVIDIA Corporation GP100GL [Tesla P100 PCIe 16GB] (rev a1)
事前準備
はじめに、必要なパッケージ群をインポートします。メインで検証するscikit-learnとcupyのバージョンは以下の通りです。
import numpy as np
import sklearn
from sklearn import config_context
from sklearn.datasets import make_classification
from sklearn.discriminant_analysis import LinearDiscriminantAnalysis
from sklearn.utils._array_api import _estimator_with_converted_arrays
import cupy as cp
import cupy.array_api as xp
for pkg in (sklearn, cp):
print(f'{pkg.__name__:<8} : {pkg.__version__}')
sklearn : 1.2.2
cupy : 11.6.0
/tmp/ipykernel_23/4001940879.py:8: UserWarning: The cupy.array_api submodule is still experimental. See NEP 47.
import cupy.array_api as xp
2クラス分類用のサンプルデータを作成します。データ件数は1000万件、X_npの特徴量数は10です。n_samplesとn_featuresの値は、検証マシンのCPU/GPUメモリサイズに応じて適時調整してください。
X_np, y_np = make_classification(n_samples=10 ** 7, n_features=10, random_state=0)
print(X_np.shape)
(10000000, 10)
y_npでは、クラス番号（0と1）が均衡に生成されています。
print(y_np.sum() / len(y_np))
0.5000255
CPUで学習
まずは従来通り、CPUとnumpyで学習・推論します。学習モデルはLinearDiscriminantAnalysisという線型判別モデルを使用します。CPUによる学習・推論の所要時間は約16秒でした。
%%time
lda_np = LinearDiscriminantAnalysis(solver='svd')
X_trans = lda_np.fit_transform(X_np, y_np)
print(f'intercept : {lda_np.intercept_}')
print(f'coef : {lda_np.coef_}')
print(f'type: {type(X_trans)}')
intercept : [0.00029512]
coef : [[ 4.13520415e-01 1.00211581e-03 1.14883736e-03 2.15778194e+00
8.53272418e-06 -2.37350983e-04 -3.96646296e-04 -2.25224889e-04
6.57963198e-01 6.55414658e-01]]
type: <class 'numpy.ndarray'>
CPU times: user 17 s, sys: 1.95 s, total: 18.9 s
Wall time: 16.3 s
GPUで学習
続いて、同じモデルをGPUとcupyで学習してみます。
まずは学習データをnumpy.ndarrayオブジェクトからcupy.ndarrayオブジェクトに変換し、データをGPUメモリに載せます。
X_cu = xp.asarray(X_np)
y_cu = xp.asarray(y_np)
print(X_cu.device, y_cu.device)
<CUDA Device 0> <CUDA Device 0>
GPU学習時には、config_context(array_api_dispatch=True)というコンテキストマネージャーの中でfit()やtransform()、predict()を実行します。
なお、GPU学習では初回のみ何らかの初期処理が行われているらしく、これだけで20秒以上の待ち時間が発生します。このままでは学習時間を正確に計測できないため、学習データの先頭10行だけで学習させ、初期処理を事前に済ませておきます。
%%time
with config_context(array_api_dispatch=True):
lda_cu = LinearDiscriminantAnalysis(solver='svd')
lda_cu.fit(X_cu[:10, :], y_cu[:10])
del lda_cu
CPU times: user 21.1 s, sys: 652 ms, total: 21.8 s
Wall time: 24.2 s
さあ、いよいよ本番です！全データを使用して学習します。
実行すると、2秒もかからずに学習が終了しました！速い！ CPU時と比較して8倍以上も高速化されています。
%%time
with config_context(array_api_dispatch=True):
lda_cu = LinearDiscriminantAnalysis(solver='svd')
X_trans = lda_cu.fit_transform(X_cu, y_cu)
CPU times: user 1.69 s, sys: 65.2 ms, total: 1.76 s
Wall time: 1.78 s
GPU学習後に得られたパラメータは、CPU学習時と同値であることが確認できます。また、transform()の出力はcupy.ndarrayオブジェクトなっており、GPUメモリに出力されます。
print(f'intercept : {lda_cu.intercept_}')
print(f'coef : {lda_cu.coef_}')
print(f'device: {X_trans.device}')
intercept : [0.00029512]
coef : [[ 4.13520415e-01 1.00211581e-03 1.14883736e-03 2.15778194e+00
8.53272419e-06 -2.37350983e-04 -3.96646296e-04 -2.25224889e-04
6.57963198e-01 6.55414658e-01]]
device: <CUDA Device 0>
GPUで学習したモデルをCPUで推論
GPUで学習したモデルは通常、推論もGPUで実行する必要があります。CPUで推論したい時は、_estimator_with_converted_arrays()というモデル変換関数を使用します。
以下のように、第1引数にGPUで学習したモデル、第2引数にcupy.ndarrayからnumpy.ndarrayへの変換関数を指定することで、CPUで推論可能な学習済みモデルが生成されます。あとは従来通り、numpy.ndarrayのデータを与えるだけでtransform()が実行できます。
converter = lambda arr : arr._array.get()
lda_cu2np = _estimator_with_converted_arrays(lda_cu, converter=converter)
X_trans = lda_cu2np.transform(X_np)
print(type(X_trans))
<class 'numpy.ndarray'>
（後始末）GPUメモリを解放
最後に、GPUメモリに載せたデータを削除します。コードによる検証はこれで以上です。
pool = cp.get_default_memory_pool()
pool.free_all_blocks()
今後、GPUで利用可能な変換器・予測器は増えていきそう？
現時点で利用可能な予測器は1種類のみですが、PRを見るとPCA(主成分分析)やMinMaxScalerなど、いくつかの前処理にてArray-API対応が提案されています。scikit-learnのコア機能はcpythonで実装されており、それを大幅に修正する必要があることを考えると「矢継ぎ早にリリースされる」ということは無さそうですが、GPUを利用可能な変換器や予測器は着実に増えてくるだろうという印象です。
まとめ
今回は、scikit-learnで実験的に導入された「GPUによる学習・推論」について検証しました。GPUによる高速演算が体感でき、この記事を書いた日はとても白熱したのを覚えています。RandomForestやHistGradientBoostingなどにも適用されれば、Kaggle等のデータ分析コンペでも大いに活躍することが期待されます。今後のアップデートが楽しみです😊