Optuna を使用したハイパーパラメーターの調整

Optuna は、複数のコンピューティング リソース間で水平方向にスケーリングできるハイパーパラメーターの調整用のオープンソース Python ライブラリです。

MLflow 3.0 では、Optuna と統合することで、ハイパーパラメーターの最適化のための強力な新機能が導入されています。

• MlflowStorage クラスを使用すると、Optuna は MLflow Tracking Server をストレージ バックエンドとして使用できます。
• MlflowSparkStudy クラスを使用すると、PySpark Executor を使用して並列 Optuna スタディを起動できます。

Optuna をインストールする

MLflow 3.0 は、Databricks Runtime 17.0 ML 以降にプレインストールされています。 古いランタイムでは、次のコマンドを使用して、Optuna と MLFlow の最新バージョンをインストールします。

%pip install mlflow --upgrade
%pip install optuna

Optuna 最適化を並列で実行する

Optuna ワークフローの手順を次に示します。

1. 最適化する目的関数を定義します。 目的関数内で、ハイパーパラメーター検索空間を定義します。 詳細については、 Optuna のドキュメントを参照してください。

   sckit-learn を使用したモデルの選択とハイパーパラメーターの調整の例を次に示します。 この例では、目的関数 objectiveを定義し、 suggest_float 関数を呼び出して、パラメーター xの検索領域を定義します。

import sklearn

def objective(trial):
    # Invoke suggest methods of a Trial object to generate hyperparameters.
    regressor_name = trial.suggest_categorical('classifier', ['SVR', 'RandomForest'])
    if regressor_name == 'SVR':
        svr_c = trial.suggest_float('svr_c', 1e-10, 1e10, log=True)
        regressor_obj = sklearn.svm.SVR(C=svr_c)
    else:
        rf_max_depth = trial.suggest_int('rf_max_depth', 2, 32)
        regressor_obj = sklearn.ensemble.RandomForestRegressor(max_depth=rf_max_depth)

    X, y = sklearn.datasets.fetch_california_housing(return_X_y=True)
    X_train, X_val, y_train, y_val = sklearn.model_selection.train_test_split(X, y, random_state=0)

    regressor_obj.fit(X_train, y_train)
    y_pred = regressor_obj.predict(X_val)

    error = sklearn.metrics.mean_squared_error(y_val, y_pred)

    return error  # An objective value linked with the Trial object

2. 分散最適化用の共有ストレージを作成します。 MlflowStorageでは、MLflow Tracking Server をストレージ バックエンドとして使用できます。

import mlflow
from mlflow.optuna.storage import MlflowStorage

experiment_id = mlflow.get_experiment_by_name(dbutils.notebook.entry_point.getDbutils().notebook().getContext().notebookPath().get()).experiment_id

mlflow_storage = MlflowStorage(experiment_id=experiment_id)

3. Optuna Study オブジェクトを作成し、Study オブジェクトの optimize 関数を呼び出してチューニング アルゴリズムを実行します。 MlflowSparkStudy PySpark Executor を使用して並列 Optuna スタディの起動を実行できます。

Optuna ドキュメントの例を次に示します。

• スタディを作成し、8 つの試行 (objective の値が異なる objective 関数の 8 回の呼び出し) でx関数を最適化します。
• Study の最適なパラメーターを取得する

from mlflow.pyspark.optuna.study import MlflowSparkStudy

mlflow_study = MlflowSparkStudy(
    study_name="spark-mlflow-tuning",
    storage=mlflow_storage,
)

mlflow_study.optimize(objective, n_trials=8, n_jobs=4)

best_params = study.best_params

ノートブックの例

このノートブックでは、Optuna を使用して Iris データセットの scikit-learn モデルと、一連のハイパーパラメーターを選択する例を示します。

Optuna と MLflow を使用してハイパーパラメーターの調整をスケールアップする

ノートブックを入手

MLFlow Optuna Integration API

MlflowStorage

MlflowStorage は、REST API の調整を回避するためのバッチ処理を備えた Optuna の MLflow ベースのストレージ クラスです。

MlflowSparkStudy

MlflowSparkStudy は、MLflow 実験を介して Optuna と Spark を組み込むクラス ~optuna.study.Study のラッパーです。