optuna_integration.WeightsAndBiasesCallback

class optuna_integration.WeightsAndBiasesCallback(metric_name='value', wandb_kwargs=None, as_multirun=False)[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


study = optuna.create_study()

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

study.optimize(objective, n_trials=10, callbacks=[wandbc])

Weights & Biases logging in multirun mode.

import optuna
from optuna_integration.wandb import WeightsAndBiasesCallback

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


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


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

Parameters:

metric_name (str | Sequence[str]) – 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. The number of metrics must be the same as the number of values objective function returns.
wandb_kwargs (dict[str, Any] | None) – Set of arguments passed when initializing Weights & Biases run. Please refer to Weights & Biases API documentation for more details.
as_multirun (bool) – Creates new runs for each trial. Useful for generating W&B Sweeps like panels (for ex., parameter importance, parallel coordinates, etc).

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.

Methods

track_in_wandb()

Decorator for using W&B for logging inside the objective function.

track_in_wandb()[source]

Decorator for using W&B for logging inside the objective function.

The run is initialized with the same wandb_kwargs that are passed to the callback. All the metrics from inside the objective function will be logged into the same run which stores the parameters for a given trial.

Example

Add additional logging to Weights & Biases.

import optuna
from optuna_integration.wandb import WeightsAndBiasesCallback
import wandb

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


@wandbc.track_in_wandb()
def objective(trial):
    x = trial.suggest_float("x", -10, 10)
    wandb.log({"power": 2, "base of metric": x - 2})

    return (x - 2) ** 2


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

Returns:: Objective function with W&B tracking enabled.
Return type:: Callable

Note

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