More than 1 year has passed since last update.

Databricksにおける機械学習モデルトレーニングのウォークスルー

Posted at 2023-11-15

以前のエンドツーエンドの機械学習のサンプルに加えて、よりシンプルなサンプルが追加されていました。

翻訳済みのノートブックはこちらです。

Databricks MLクイックスタート: モデルのトレーニング

このノートブックでは、Databricksにおける機械学習モデルトレーニングの概要を説明します。モデルをトレーニングするには、Databricks機械学習ランタイムにプレインストールされているscikit-learnのようなライブラリを活用することができます。さらに、トレーニングしたモデルの追跡にMLflowを活用し、ハイパーパラメータチューニングをスケールさせるためにSparkTrialsとHyperoptを活用することもできます。

このチュートリアルでは以下をカバーします:

Part 1: MLflowトラッキングを用いたシンプルな分類モデルのトレーニング
Part 2: Hyperoptを用いたより性能の良いモデルをトレーニングするためのハイパーパラメーターチューニング

モデルライフサイクル管理とモデル推論を含む、Databricksにおける機械学習のプロダクション化の詳細に関しては、MLのエンドツーエンドのサンプルをご覧ください (AWS|Azure|GCP)。

要件

Databricksランタイム7.5 ML以降が稼働しているクラスター

ライブラリ

必要なライブラリをインポートします。これらのライブラリは、Databricks機械学習ランタイム(AWS|Azure|GCP)クラスターにプレインストールされており、互換性とパフォーマンスにチューニングされています。

Python

import mlflow
import numpy as np
import pandas as pd
import sklearn.datasets
import sklearn.metrics
import sklearn.model_selection
import sklearn.ensemble

from hyperopt import fmin, tpe, hp, SparkTrials, Trials, STATUS_OK
from hyperopt.pyll import scope

データのロード

このチュートリアルでは、様々なワインのサンプルを記述しているデータセットを使用します。このデータセットはUCI機械学習リポジトリから取得しているものであり、DBFS (AWS|Azure|GCP)に格納されています。赤ワインと白ワインを品質に基づいて分類することがゴールとなります。

アップロードや他のデータソースからのロードに関する詳細については、データ取扱に関するドキュメント(AWS|Azure|GCP)をご覧ください。

Python

# データのロードと前処理
white_wine = pd.read_csv("/dbfs/databricks-datasets/wine-quality/winequality-white.csv", sep=';')
red_wine = pd.read_csv("/dbfs/databricks-datasets/wine-quality/winequality-red.csv", sep=';')
white_wine['is_red'] = 0.0
red_wine['is_red'] = 1.0
data_df = pd.concat([white_wine, red_wine], axis=0)

# ワイン品質に基づいた分類ラベルの定義
data_labels = data_df['quality'] >= 7
data_df = data_df.drop(['quality'], axis=1)

# 80/20でトレーニング/テストデータセットを分割
X_train, X_test, y_train, y_test = sklearn.model_selection.train_test_split(
  data_df,
  data_labels,
  test_size=0.2,
  random_state=1
)

Part 1. 分類モデルのトレーニング

MLflowトラッキング

MLflowトラッキングを用いることで、お使いの機械学習トレーニングコード、パラメータ、モデルを整理することができます。

autologgingを用いることで、自動でのMLflowトラッキングを有効化することができます。

Python

# このノートブックでのMLflow autologgingを有効化
mlflow.autolog()

2023/11/15 03:18:38 INFO mlflow.tracking.fluent: Autologging successfully enabled for sklearn.
2023/11/15 03:18:39 INFO mlflow.tracking.fluent: Autologging successfully enabled for pyspark.
2023/11/15 03:18:39 INFO mlflow.tracking.fluent: Autologging successfully enabled for pyspark.ml.

次に、トレーニングされたモデルや関連付けられるメトリクスやパラメータが自動で記録されるMLflowランのコンテキスト内で分類器をトレーニングします。

テストデータに対するモデルのAUCスコアのようなその他のメトリクスを追加することも可能です。

Python

with mlflow.start_run(run_name='gradient_boost') as run:
  model = sklearn.ensemble.GradientBoostingClassifier(random_state=0)
  
  # モデル、パラメータ、トレーニングメトリクスが自動でトラッキングされます
  model.fit(X_train, y_train)

  predicted_probs = model.predict_proba(X_test)
  roc_auc = sklearn.metrics.roc_auc_score(y_test, predicted_probs[:,1])
  
  # テストデータに対するAUCスコアは自動で記録されないので、手動で記録します
  mlflow.log_metric("test_auc", roc_auc)
  print("Test AUC of: {}".format(roc_auc))

画面右のフラスコマークをクリックすると記録されたモデルを確認することができます。

このモデルのパフォーマンスに満足しない場合には、異なるハイパーパラメーターを用いて別のモデルをトレーニングします。

Python

# 新たなランをスタートし、後でわかるようにrun_nameを割り当てます
with mlflow.start_run(run_name='gradient_boost') as run:
  model_2 = sklearn.ensemble.GradientBoostingClassifier(
    random_state=0, 
    
    # n_estimatorsで新たなパラメータ設定をトライします
    n_estimators=200,
  )
  model_2.fit(X_train, y_train)

  predicted_probs = model_2.predict_proba(X_test)
  roc_auc = sklearn.metrics.roc_auc_score(y_test, predicted_probs[:,1])
  mlflow.log_metric("test_auc", roc_auc)
  print("Test AUC of: {}".format(roc_auc))

ハイパーパラメーターが異なる2つ目のモデルが記録されます。

MLflowランの参照

記録されたトレーニングランを参照するには、エクスペリメントサイドバーを表示するためにノートブック右上のExperimentアイコンをクリックします。必要であれば、最新のランを取得、監視するためにリフレッシュアイコンをクリックします。

より詳細なMLflowエクスペリメントページ(AWS|Azure|GCP)を表示するために、エクスペリメントページアイコンをクリックすることもできます。このページでは、複数のランを比較し、特定のランの詳細を表示することができます。

モデルのロード

また、MLflow APIを用いて特定のランの結果にアクセスすることができます。以下のセルのコードでは、特定のMLflowランでトレーニングされたモデルのロード方法と、予測での利用方法を説明しています。また、MLflowランのページ(AWS|Azure|GCP)では特定のモデルをロードするためのコードスニペットを参照することができます。

Python

# モデルが記録された後では、別のノートブックやジョブでモデルをロードすることができます
# mlflow.pyfunc.load_modelを用いることで、共通APIを通じてモデルの予測を行うことができます
model_loaded = mlflow.pyfunc.load_model(
  'runs:/{run_id}/model'.format(
    run_id=run.info.run_id
  )
)

predictions_loaded = model_loaded.predict(X_test)
predictions_original = model_2.predict(X_test)

# ロードされたモデルはオリジナルと一致すべきです
assert(np.array_equal(predictions_loaded, predictions_original))

結果が一致するのでassertionを通過します。

Part 2. ハイパーパラメーターチューニング

この時点で、シンプルなモデルをトレーニングし、皆様の取り組み結果を整理するためにMLflowトラッキングサービスを活用しました。このセクションでは、Hyperoptを用いてより洗練されたチューニングを実行する方法をカバーします。

HyperoptとSparkTrialsによる並列トレーニング

HyperoptはハイパーパラメーターチューニングのためのPythonライブラリです。DatabricksにおけるHyperopt活用の詳細については、ドキュメント(AWS|Azure|GCP)を参照ください。

並列でのハイパーパラメーター探索と複数のモデルのトレーニング実行を行うために、HyperoptとSparkTrialsを活用することができます。これによって、モデルパフォーマンスの最適化に要する時間を削減します。モデルとパラメータを自動で記録するように、MLflowトラッキングはHyperoptとインテグレーションされています。

Python

# 探索する検索空間の定義
search_space = {
  'n_estimators': scope.int(hp.quniform('n_estimators', 20, 1000, 1)),
  'learning_rate': hp.loguniform('learning_rate', -3, 0),
  'max_depth': scope.int(hp.quniform('max_depth', 2, 5, 1)),
}

def train_model(params):
  # それぞれのワーカーでautologgingを有効化
  mlflow.autolog()
  with mlflow.start_run(nested=True):
    model_hp = sklearn.ensemble.GradientBoostingClassifier(
      random_state=0,
      **params
    )
    model_hp.fit(X_train, y_train)
    predicted_probs = model_hp.predict_proba(X_test)
    # テストデータに対するAUCに基づくチューニング
    # プロダクション環境では、代わりに別の検証用データセットを活用することができます
    roc_auc = sklearn.metrics.roc_auc_score(y_test, predicted_probs[:,1])
    mlflow.log_metric('test_auc', roc_auc)
    
    # fminがauc_scoreを最大化するように、lossを -1*auc_score に設定します
    return {'status': STATUS_OK, 'loss': -1*roc_auc}

# SparkTrialsがSparkワーカーを用いてチューニングを分散させます
# 並列度を高めると処理を加速しますが、それぞれのハイパーパラメーターのトライアルが他のトライアルから得られる情報が減少します
# 小規模なクラスターやDatabricks Community Editionでは parallelism=2 に設定してください
spark_trials = SparkTrials(
  parallelism=8
)

with mlflow.start_run(run_name='gb_hyperopt') as run:
  # 最大のAUCを達成するパラメーターを特定するためhyperoptを使用します
  best_params = fmin(
    fn=train_model, 
    space=search_space, 
    algo=tpe.suggest, 
    max_evals=32,
    trials=spark_trials)

ハイパーパラメーターチューニングの過程で多数のモデルが生成されますが、それらはすべてMLflowによって追跡・記録されます。

ベストモデルを取得するためにランを検索します

すべてのランはMLflowで追跡されるので、テストデータに対する最大のAUCを持つチューニングのランを特定するために、MLflowのsearch runs APIを用いて、ベストなランのメトリクスとパラメーターを取得することができます。

このチューニングされたモデルは、パート1でトレーニングされたよりシンプルなモデルよりも良いパフォーマンスであるべきです。

Python

# テストデータに対するAUCでランをソートします。値が同じ場合には最新のランを使用します
best_run = mlflow.search_runs(
  order_by=['metrics.test_auc DESC', 'start_time DESC'],
  max_results=10,
).iloc[0]
print('Best Run')
print('AUC: {}'.format(best_run["metrics.test_auc"]))
print('Num Estimators: {}'.format(best_run["params.n_estimators"]))
print('Max Depth: {}'.format(best_run["params.max_depth"]))
print('Learning Rate: {}'.format(best_run["params.learning_rate"]))

best_model_pyfunc = mlflow.pyfunc.load_model(
  'runs:/{run_id}/model'.format(
    run_id=best_run.run_id
  )
)
best_model_predictions = best_model_pyfunc.predict(X_test[:5])
print("Test Predictions: {}".format(best_model_predictions))

ハイパーパラメーターチューニングによって、AUC(Area Under Curve)は0.91から0.917まで改善しました。

Best Run
AUC: 0.9174331940947587
Num Estimators: 935.0
Max Depth: 5.0
Learning Rate: 0.09614045842606779

UIで複数のランを比較

パート1と同じように、Experimentサイドバーの上部にある外部リンクアイコン経由でMLflowエクスペリメントの詳細ページでランを参照、比較できます。

エクスペリメント詳細ページでは、親のランを展開するために　"+"　アイコンをクリックし、親を除くすべてのランを選択し、Compareをクリックします。メトリックに対する様々なパラメーターの値のインパクトを表示するparallel coordinates plotを用いて、様々なランを可視化することができます。

まとめ

世の中は生成AI一色ではありますが、世の中のユースケースすべてが生成AIという訳ではありません。この例で触れたような、構造化データの特徴量に基づく機械学習モデルのニーズも依然として存在しており、そのようなモデル開発・運用の営みはこれまで以上に増えてくると思います。そういった場合に、MLflowなどによって実現されるMLOpsの重要性も同様に高まってくると考えております。是非、DatabricksでMLflowを活用いただき、MLOpsを実現してください！

Databricksクイックスタートガイド

Databricks無料トライアル

You get articles that match your needs
You can efficiently read back useful information
You can use dark theme

What you can do with signing up