optuna.integration.MLflowCallback¶

class optuna.integration.MLflowCallback(tracking_uri=None, metric_name='value', nest_trials=False, tag_study_user_attrs=False)[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_float("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])

Add additional logging to MLflow

import optuna
import mlflow
from optuna.integration.mlflow import MLflowCallback

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


@mlflc.track_in_mlflow()
def objective(trial):
    x = trial.suggest_float("x", -10, 10)
    mlflow.log_param("power", 2)
    mlflow.log_metric("base of metric", x - 2)

    return (x - 2) ** 2


study = optuna.create_study(study_name="my_other_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 assigned to optimized metric. In case of multi-objective optimization, list of names can be passed. Those names will be assigned to metrics in the order returned by objective function. If single name is provided, or this argument is left to default value, it will be broadcasted to each objective with a number suffix in order returned by objective function e.g. two objectives and default metric name will be logged as value_0 and value_1.
nest_trials – Flag indicating whether or not trials should be logged as nested runs. This is often helpful for aggregating trials to a particular study, under a given experiment.
tag_study_user_attrs – Flag indicating whether or not to add the study’s user attrs to the mlflow trial as tags. Please note that when this flag is set, key value pairs in user_attrs will supersede existing tags.

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.

Methods

track_in_mlflow()

Decorator for using MLflow logging in the objective function.

track_in_mlflow()[source]¶

Decorator for using MLflow logging in the objective function.

This decorator enables the extension of MLflow logging provided by the callback.

All information logged in the decorated objective function will be added to the MLflow run for the trial created by the callback.

Returns: Objective function with tracking to MLflow enabled.
Return type: ObjectiveFuncType

Note

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