Scikit-learnが実験的にGPU対応していたので調査してみた！

## はじめに
みずほリサーチ＆テクノロジーズ株式会社の@fujineです。

いきなりですがAIエンジニアの皆さん、**scikit-learnが実験的にGPUに対応していた**こと、ご存知でしょうか？

scikit-learnは機械学習分野における古参パッケージの1つです。多様な機能を提供する一方、[FAQにて「GPUに対応する予定はない（キリッ）」](https://scikit-learn.org/stable/faq.html#will-you-add-gpu-support)と公式宣言しており、scikit-learnが好きな自分としては「勿体無いなぁ」と常々感じていました。

そんな中、何気なく[Release Highlights for 1.2](https://scikit-learn.org/stable/auto_examples/release_highlights/plot_release_highlights_1_2_0.html#experimental-array-api-support-in-lineardiscriminantanalysis)を読んでいたら以下文面を発見！しかも約半年前の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](https://scikit-learn.org/stable/modules/array_api.html#example-usage)に沿って検証しました。

## 概要
- [scikit-learnのver1.2](https://scikit-learn.org/stable/auto_examples/release_highlights/plot_release_highlights_1_2_0.html)にて、[GPUによる学習・推論が実験的に導入](https://scikit-learn.org/stable/modules/array_api.html#array-api-support-experimental)された
- 現時点では、まだ1種類の線型判別器（[LinearDiscriminantAnalysis](https://scikit-learn.org/stable/modules/generated/sklearn.discriminant_analysis.LinearDiscriminantAnalysis.html)）のみに適用可能
- GitHubでは[issue](https://github.com/scikit-learn/scikit-learn/issues?q=is%3Aissue+label%3A%22Array+API%22+)や[PR](https://github.com/scikit-learn/scikit-learn/pulls?q=is%3Apr+label%3A%22Array+API%22+)が活発に提案されており、今後様々な前処理や予測器がGPU上で実行可能になると期待される

## 検証環境
お手軽にKaggleのNotebook環境を利用します。オプションにて、ACCELERATORは[Tesla P100](https://www.nvidia.com/ja-jp/data-center/tesla-p100/)、ENVIRONMENTは「Always use latest environment」をそれぞれ選択します。

![env.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/2969456/b0b203e1-49ae-cdfd-65f2-0b2f13053065.png)

検証環境のPythonとOSのバージョン、GPUデバイスの情報は以下の通りです。


```sh
$ 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`のバージョンは以下の通りです。


```python
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メモリサイズに応じて適時調整してください。


```python
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`）が均衡に生成されています。


```python
print(y_np.sum() / len(y_np))
```

    0.5000255
    

## CPUで学習
まずは従来通り、CPUと`numpy`で学習・推論します。学習モデルは[LinearDiscriminantAnalysis](https://scikit-learn.org/stable/modules/generated/sklearn.discriminant_analysis.LinearDiscriminantAnalysis.html)という線型判別モデルを使用します。CPUによる学習・推論の所要時間は約16秒でした。


```python
%%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](https://docs.cupy.dev/en/stable/overview.html)で学習してみます。

まずは学習データを`numpy.ndarray`オブジェクトから`cupy.ndarray`オブジェクトに変換し、データをGPUメモリに載せます。


```python
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行だけで学習させ、初期処理を事前に済ませておきます。


```python
%%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倍以上も高速化**されています。


```python
%%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メモリに出力されます。


```python
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()`が実行できます。


```python
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メモリに載せたデータを削除します。コードによる検証はこれで以上です。


```python
pool = cp.get_default_memory_pool()
pool.free_all_blocks()
```

## 今後、GPUで利用可能な変換器・予測器は増えていきそう？
現時点で利用可能な予測器は1種類のみですが、PRを見ると[PCA(主成分分析)](https://github.com/scikit-learn/scikit-learn/pull/26315)や[MinMaxScaler](https://github.com/scikit-learn/scikit-learn/pull/26243)など、いくつかの前処理にてArray-API対応が提案されています。scikit-learnのコア機能はcpythonで実装されており、それを大幅に修正する必要があることを考えると「矢継ぎ早にリリースされる」ということは無さそうですが、GPUを利用可能な変換器や予測器は着実に増えてくるだろうという印象です。

## まとめ
今回は、scikit-learnで実験的に導入された「GPUによる学習・推論」について検証しました。GPUによる高速演算が体感でき、この記事を書いた日はとても白熱したのを覚えています。[RandomForest](https://scikit-learn.org/stable/modules/ensemble.html#forests-of-randomized-trees)や[HistGradientBoosting](https://scikit-learn.org/stable/modules/ensemble.html#histogram-based-gradient-boosting)などにも適用されれば、Kaggle等のデータ分析コンペでも大いに活躍することが期待されます。今後のアップデートが楽しみです😊