使用 Optuna 进行超参数优化

Optuna 是一个用于超参数优化的开源 Python 库，可在多个计算资源之间进行水平缩放。

MLflow 3.0 通过与 Optuna 集成，为超参数优化引入了强大的新功能。

• MlflowStorage 类允许 Optuna 使用 MLflow 跟踪服务器作为其存储后端。
• MlflowSparkStudy 类支持使用 PySpark 执行程序启动并行 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 跟踪服务器作为存储后端。

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 研究”对象，并通过调用该“研究”对象的 optimize 函数以运行优化算法。 MlflowSparkStudy 可以使用 PySpark 执行器运行并行 Optuna 实验。

下面是 Optuna 文档中的一个示例。

• 创建一项研究，并通过8次实验来优化objective函数（x 8 次对 objective 函数进行不同值的调用）。
• 获取“研究”的最佳参数

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 集成 API

MlflowStorage

MlflowStorage 是一个基于 MLflow 的存储类，通过批处理来避免 REST API 限流。

MlflowSparkStudy

MlflowSparkStudy 是类 ~optuna.study.Study 的包装器，用于通过 MLflow 试验将 Optuna 与 Spark 合并。