🌏

Databricks上でscikit-learnを高速化する ~シングルノード編~

2023/10/12に公開

scikit-learn

Databricks上でscikit-learnを高速化する　~シングルノード編~

Databricksには多くのデータ分析用ライブラリーが事前にインストールされており、どなたでも自在に高度な分析を実施いただけます。したがって、Databricksを分析ツールとして使用されている方は多いと思います。データ分析の過程で機械学習モデルを作るケースも多く、代表的なライブラリーとしてscikit-learnがよく使用されてます。当然、Databricksにもscikit-learnが事前インストールされています。
本ブログでは、Databricks上で誰でも簡単にscikit-learnのモデル学習処理を高速化するためのTipsの一つとして、Intel(R) Extension for Scikit-learn*（以降、Intel scikit-learnと表記）を使用する方法を紹介します。

Intel(R) Extension for Scikit-learn* とは

その名の通り、インテルがインテル CPUとインテル GPU上でのアルゴリズム実行の高速化のために最適化を行なったScikit-learnモジュールです。これ単体で動くわけではなく、オリジナルのscikit-learnのアドオンモジュールという形で動作します。したがって、前提として、オリジナルがインストールされている必要があります。

環境

本ブログではAzure Databricksを用いていますが、AWS版、GCP版のいずれでも再現可能かと思います。

Databricks Runtime: 13.3 LTS ML (includes Apache Spark 3.4.1, Scala 2.12, scikit-learn 1.1.1 pre-installed)
ノードタイプ: Azure Standard Ds_V5シリーズ（シングルノードで起動）
主要ライブラリ: Intel(R) Extension for Scikit-learn*

制約

Intel(R) Extension for Scikit-learn*ではscikit-learnのすべてのアルゴリズムが最適化されているわけではありません。こちらのURL に記載されているアルゴリズムのみ、性能向上が期待できます。今回は「RandomForestClassifier」と「LogisticRegression」を試していきます。
また、原則インテル社のCPUおよびGPUベースのインスタンスのみサポートされます。（２０２３年１０月１２日現在、AWS／Azure／GCPではインテルGPUインスタンスというのは存在しないので、実質CPUインスタンスのみが￥対象）。AMD CPU、または、NVIDIA GPUのインスタンスでは動作が保証されていませんのでご注意を。

コード

ベースとしてこちらの公式サンプルを利用しました。
その上で、以下の点を追加・変更しました。

Intel(R) Extension for Scikit-learn*をインストール

%pip install scikit-learn-intelex

Pythonランタイムの再起動

dbutils.library.restartPython()

Intel(R) Extension for Scikit-learn*のインポートおよびアクティベーション追加

from sklearnex import patch_sklearn
patch_sklearn()

データ拡張
元のノートブックで使用しているデータのレコード数が６,４９７だけなので、最適化前でも一瞬で学習が終了してしまい、性能差が出づらいと考え、レコード数を１００倍に増幅しました。

def data_augmentation(df, scale : int):
  new_df = None
  for i in range(scale):
    new_df = pd.concat([new_df, df], axis=0)

  print(f"The data is successfully augumented from {df.shape} to {new_df.shape}")
  return new_df

#レコード数に１００倍の増幅
data_df = data_augmentation(data_df, 100)

使用するアルゴリズムを元の「GradientBoostingClassifier」から「RandomForestClassifier」に変更
ちなみにRandomForestにはn_jobsというパラメータがあり、その名の通りジョブ数を指定できます。ジョブ数が増えるほど処理の並列度が増し、性能向上が期待できます。が、一旦このパラメータはデフォルト値（つまりジョブ数は１）で実行します。n_jobsについてはこちら

with mlflow.start_run(run_name='random_forest') as run:
  #model = sklearn.ensemble.GradientBoostingClassifier(random_state=0)
  model = sklearn.ensemble.RandomForestClassifier(random_state=0)
  
  # Models, parameters, and training metrics are tracked automatically
  start = time.time()
  model.fit(X_train, y_train)
  print(time.time() - start)

  predicted_probs = model.predict_proba(X_test)
  roc_auc = sklearn.metrics.roc_auc_score(y_test, predicted_probs[:,1])
  
  # The AUC score on test data is not automatically logged, so log it manually
  mlflow.log_metric("test_auc", roc_auc)
  print("Test AUC of: {}".format(roc_auc))

RandomForest(n_jobs=1)の学習処理結果

以下が結果です。２種類のノードタイプ（Azure VM）で検証しています。

ノードタイプ	vCPU	物理コア	scikit-learnのみ (秒)	scikit-learn + Intel scikit-learn (秒)
Standard_D4s_v5	4	2	69.60	28.73
Standard_D16s_v5	16	8	67.58	11.76

https://learn.microsoft.com/ja-jp/azure/virtual-machines/dv5-dsv5-series

VMサイズが大きくなる、つまり、物理コアの数が増えると、性能向上の割合が上がることが確認できました。これは、ジョブ数は１にも関わらず、オリジナルのScikit-learnと比較して、Intel Scikit-learnがマルチコアを効果的に利用するためのマルチスレッド処理がより最適化されているという証拠と言えるでしょう。逆にオリジナルのScikit-learnでは、コア数が増えても性能に大きな変化が見られないため、マルチコアをほとんど利用していない（基本的にシングルスレッド動作か？）と考えられます。

RandomForest(n_jobs=物理コア数と同一数)の学習処理結果

せっかくなのでn_jobsをコア数と同一の数値にしてみましょう。この条件でもIntel scikit-learnがどの程度性能向上するか確認していきます。


  #全コアを使う場合は n_jobs=-1 と設定します
  model = sklearn.ensemble.RandomForestClassifier(random_state=0, n_jobs=-1)

ノードタイプ	vCPU	物理コア	scikit-learnのみ (秒)	scikit-learn + Intel scikit-learn (秒)
Standard_D4s_v5	4	2	33.18	28.57
Standard_D16s_v5	16	8	14.57	11.54

流石に並列度を上げるとオリジナルscikit-learnでもそれなりに性能が上がりました。ただ、Intel scikit-learnの方がまだ１.２倍ほど速いことがわかるかと思います。公式ベンチマークでも、RandomForestは1.4倍ほどの性能向上と謳っているので、上記の結果は妥当かと思います。

Logistic Regressionの学習処理結果

別のアルゴリズムでも試してみましょう。Logistic Regressionで行きます。


  #全コアを使う場合は n_jobs=-1 と設定します。
  #また、モデルの収束する前に反復が終了してしまう
  #  「STOP: TOTAL NO. of ITERATIONS REACHED LIMIT.」が
  #  出るのを防ぐために、max_iterに十分大きな数字（とりあえず10000）をセットしています。
  model = sklearn.linear_model.LogisticRegression(random_state=0, max_iter=10000, n_jobs=-1)

ノードタイプ	vCPU	物理コア	scikit-learnのみ (秒)	scikit-learn + Intel scikit-learn (秒)
Standard_D4s_v5	4	2	32.62	14.28
Standard_D16s_v5	16	8	30.86	6.52

先ほどのRandomForestと比べると、n_jobsを設定しているにも関わらず、かなり顕著な性能向上が確認できますね。またVMサイズの違い（=コア数の違い）によって、性能向上の割合が増えている点も興味深いです。
このようにアルゴリズムによって性能の伸び方は大きく異なります。Logistic Regressionは１層のニューラルネットワークと考えることができ、お馴染みの行列計算が多くを占めているので、CPU内のSIMD（Single instruction, multiple data）命令、および、コア数が生きてくるアルゴリズムとなります。

今回使用したAzure VM DS v5シリーズに搭載されているCPUは「第3世代 Intel® Xeon® Platinum 8370C」で、SIMDの拡張命令セット「Intel® Advanced-Vector Extensions 512 (Intel® AVX-512)」をサポートしています。Intel Scikit-learnは、このAVX-512を最大限に活用するための最適化が施されており、それが上記の結果に繋がっています。興味のある方は、Githubで詳細なコードを確認してみてください。

まとめ

いずれにしても、VMサイズを変えずとも、ライブラリを一つ追加するだけで、いずれのアルゴリズムも性能向上が実現できたので、TCO削減に向けた有効なオプションの一つとして活用できるかと思います。（２物理コアのVMでも性能が上がったのは嬉しい誤算でした）
ぜひ皆様のDatabricks環境でもお試しください。

ただ、何度も言いますが、性能向上が期待できるアルゴリズムは限定的なので、公式ドキュメントを参照し、自身が使用しているアルゴリズムがリストに含まれているかをご確認ください。

こういった最適化が効いてくるのは、今回のような数十秒程度のワークロードよりも、ハイパーパラメータチューニングのような長時間ワークロードですよね。というわけで次回は今回の延長として、Hyperopt と SparkTrialsを使用して、Sparkクラスタ上でハイパーパラメーターチューニングの処理の高速化を検証する予定です。

BFN!

参考

Use scikit-learn on Databricks
Intel(R) Extension for Scikit-learn*