Ajuste de hiperparâmetros com o Optuna

O Optuna é uma biblioteca Python de código aberto para ajuste de hiperparâmetros que pode ser dimensionada horizontalmente em vários recursos de computação.

O MLflow 3.0 introduz novos e poderosos recursos para otimização de hiperparâmetros através da integração com o Optuna.

• MlflowStorage permite que o Optuna use o MLflow Tracking Server como seu back-end de armazenamento.
• MlflowSparkStudy classe permite iniciar estudos paralelos do Optuna usando executores PySpark.

Instalar o Optuna

O MLflow 3.0 está pré-instalado no Databricks Runtime 17.0 ML e superior. Em tempos de execução mais antigos, use os seguintes comandos para instalar a versão mais recente do Optuna e MLFlow.

%pip install mlflow --upgrade
%pip install optuna

Execute a otimização do Optuna em paralelo

Aqui estão as etapas em um fluxo de trabalho do Optuna:

1. Defina uma função objetiva para otimizar. Dentro da função objetivo, defina o espaço de pesquisa de hiperparâmetros. Para obter mais detalhes, consulte a documentação do Optuna.

   Abaixo está um exemplo de seleção de modelos e ajuste de hiperparâmetros com sckit-learn. O exemplo define a função objectiveobjetive e chama a suggest_float função para definir o espaço de pesquisa para o parâmetro 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. Crie um armazenamento compartilhado para otimização distribuída. Com MlflowStorage, você pode usar o MLflow Tracking Server como back-end de armazenamento.

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. Crie um objeto Optuna Study e execute o algoritmo de ajuste chamando a optimize função do objeto Study. MlflowSparkStudy pode executar o lançamento de estudos paralelos do Optuna usando executores PySpark.

Abaixo está um exemplo da documentação do Optuna.

• Crie um estudo e otimize a objective função com 8 tentativas (8 chamadas da objective função com diferentes valores de x).
• Obtenha os melhores parâmetros do Estudo

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

Exemplo de bloco de notas

Este bloco de anotações fornece um exemplo de uso do Optuna para selecionar um modelo scikit-learn e um conjunto de hiperparâmetros para o conjunto de dados Iris.

Ampliando o ajuste de hiperparâmetros com Optuna e MLflow

Obter o bloco de notas

API de integração MLFlow Optuna

MlflowStorage

MlflowStorage é uma classe de armazenamento baseada em MLflow para o Optuna com processamento em lote para evitar as limitações da API REST.

Estudo MlflowSparkStudy

MlflowSparkStudy é um wrapper da classe ~optuna.study.Study para incorporar o Optuna com o Spark via experimento MLflow.