optuna.integration.WeightsAndBiasesCallback

class optuna.integration.WeightsAndBiasesCallback(metric_name='value', wandb_kwargs=None)[source]

Callback to track Optuna trials with Weights & Biases.

This callback enables tracking of Optuna study in Weights & Biases. The study is tracked as a single experiment run, where all suggested hyperparameters and optimized metrics are logged and plotted as a function of optimizer steps.

Note

User needs to be logged in to Weights & Biases before using this callback in online mode. For more information, please refer to wandb setup.

Note

Users who want to run multiple Optuna studies within the same process should call wandb.finish() between subsequent calls to study.optimize(). Calling wandb.finish() is not necessary if you are running one Optuna study per process.

Note

To ensure correct trial order in Weights & Biases, this callback should only be used with study.optimize(n_jobs=1).

Example

Add Weights & Biases callback to Optuna optimization.

import optuna
from optuna.integration.wandb import WeightsAndBiasesCallback


def objective(trial):
    x = trial.suggest_float("x", -10, 10)
    return (x - 2) ** 2


wandb_kwargs = {"project": "my-project"}
wandbc = WeightsAndBiasesCallback(wandb_kwargs=wandb_kwargs)


study = optuna.create_study(study_name="my_study")
study.optimize(objective, n_trials=10, callbacks=[wandbc])
Parameters

  • 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.

  • wandb_kwargs – Set of arguments passed when initializing Weights & Biases run. Please refer to Weights & Biases API documentation for more details.

Raises

  • ValueError – If there are missing or extra metric names in multi-objective optimization.

  • TypeError – When metric names are not passed as sequence.

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.