optuna.integration.MLflowCallback

class optuna.integration.MLflowCallback(tracking_uri: Optional[str] = None, metric_name: str = 'value')[source]

Callback to track Optuna trials with MLflow.

This callback adds relevant information that is tracked by Optuna to MLflow. The MLflow experiment will be named after the Optuna study name.

Example

Add MLflow callback to Optuna optimization.

import optuna
from optuna.integration.mlflow import MLflowCallback

def objective(trial):
    x = trial.suggest_uniform('x', -10, 10)
    return (x - 2) ** 2

mlflc = MLflowCallback(
    tracking_uri=YOUR_TRACKING_URI,
    metric_name='my metric score',
)

study = optuna.create_study(study_name='my_study')
study.optimize(objective, n_trials=10, callbacks=[mlflc])
Parameters
  • tracking_uri

    The URI of the MLflow tracking server.

    Please refer to mlflow.set_tracking_uri for more details.

  • metric_name – Name of the metric. Since the metric itself is just a number, metric_name can be used to give it a name. So you know later if it was roc-auc or accuracy.

Note

Added in v1.4.0 as an experimental feature. The interface may change in newer versions without prior notice. See https://github.com/optuna/optuna/releases/tag/v1.4.0.

__init__(tracking_uri: Optional[str] = None, metric_name: str = 'value') None[source]

Methods

__init__([tracking_uri, metric_name])