Repository: optuna/optuna

Issue:
Weights and biases integration
<!-- Please write a clear and concise description of the feature proposal. -->

NOTE: I haven't used any management tools for ML; so it is an idea to enlarge optuna community.

## Motivation

[weights and biases](https://wandb.ai/site) (W&B) is a tool for tracking machine learning experiments, visualize metrics, and share results like MLflow. As a counterpart of MLflow, it might be great to support W&B integration for W&B users.

## Description

Like optuna's mlflow integration, optuna might provide a callback for W&B and its example.

## Alternatives (optional)

<!-- Please write a clear and concise description of any alternative solutions or features you've considered. -->

## Additional context (optional)

<!-- Please add any other context or screenshots about the feature request here. -->

[A blog post to use optuna and W&B](https://www.h4pz.co/blog/2020/10/3/optuna-and-wandb)


Discussion:
This issue has not seen any recent activity.
Hi, can I implement this callback?
@xadrianzetx  Yes! Thank you for your proposal.

Resolving patch:
diff --git a/optuna/integration/__init__.py b/optuna/integration/__init__.py
--- a/optuna/integration/__init__.py
+++ b/optuna/integration/__init__.py
@@ -13,6 +13,7 @@
     "chainermn": ["ChainerMNStudy"],
     "cma": ["CmaEsSampler", "PyCmaSampler"],
     "mlflow": ["MLflowCallback"],
+    "wandb": ["WeightsAndBiasesCallback"],
     "keras": ["KerasPruningCallback"],
     "lightgbm": ["LightGBMPruningCallback", "LightGBMTuner", "LightGBMTunerCV"],
     "pytorch_distributed": ["TorchDistributedTrial"],
@@ -61,6 +62,7 @@
     from optuna.integration.tensorboard import TensorBoardCallback  # NOQA
     from optuna.integration.tensorflow import TensorFlowPruningHook  # NOQA
     from optuna.integration.tfkeras import TFKerasPruningCallback  # NOQA
+    from optuna.integration.wandb import WeightsAndBiasesCallback  # NOQA
     from optuna.integration.xgboost import XGBoostPruningCallback  # NOQA
 else:
 
diff --git a/optuna/integration/wandb.py b/optuna/integration/wandb.py
new file mode 100644
--- /dev/null
+++ b/optuna/integration/wandb.py
@@ -0,0 +1,96 @@
+from typing import Any
+from typing import Dict
+from typing import Optional
+
+import optuna
+from optuna._experimental import experimental
+from optuna._imports import try_import
+
+
+with try_import() as _imports:
+    import wandb
+
+
+@experimental("2.9.0")
+class WeightsAndBiasesCallback(object):
+    """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 metric
+    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 <https://docs.wandb.ai/quickstart#1-set-up-wandb>`_.
+
+    .. 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)``.
+
+    .. note::
+        This callback currently supports only single-objective optimization.
+
+    Example:
+
+        Add Weights & Biases callback to Optuna optimization.
+
+        .. code::
+
+            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])
+
+    Args:
+        metric_name:
+            Name of the optimized 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.
+        wandb_kwargs:
+            Set of arguments passed when initializing Weights & Biases run.
+            Please refer to `Weights & Biases API documentation
+            <https://docs.wandb.ai/ref/python/init>`_ for more details.
+    """
+
+    def __init__(
+        self, metric_name: str = "value", wandb_kwargs: Optional[Dict[str, Any]] = None
+    ) -> None:
+
+        _imports.check()
+
+        self._metric_name = metric_name
+        self._wandb_kwargs = wandb_kwargs or {}
+
+        self._initialize_run()
+
+    def __call__(self, study: optuna.study.Study, trial: optuna.trial.FrozenTrial) -> None:
+
+        direction = study.direction.name
+        attributes = {"direction": direction}
+
+        wandb.config.update(attributes)
+        wandb.log({**trial.params, self._metric_name: trial.value}, step=trial.number)
+
+    def _initialize_run(self) -> None:
+        """Initializes Weights & Biases run."""
+
+        wandb.init(**self._wandb_kwargs)
diff --git a/setup.py b/setup.py
--- a/setup.py
+++ b/setup.py
@@ -138,6 +138,7 @@ def get_extras_require() -> Dict[str, List[str]]:
             "cma",
             "lightgbm",
             "mlflow",
+            "wandb",
             "mpi4py",
             "mxnet",
             "pandas",


Test patch:
diff --git a/tests/integration_tests/test_wandb.py b/tests/integration_tests/test_wandb.py
new file mode 100644
--- /dev/null
+++ b/tests/integration_tests/test_wandb.py
@@ -0,0 +1,73 @@
+from unittest import mock
+
+import optuna
+from optuna.integration import WeightsAndBiasesCallback
+
+
+def _objective_func(trial: optuna.trial.Trial) -> float:
+
+    x = trial.suggest_uniform("x", low=-10, high=10)
+    y = trial.suggest_loguniform("y", low=1, high=10)
+    return (x - 2) ** 2 + (y - 25) ** 2
+
+
+@mock.patch("optuna.integration.wandb.wandb")
+def test_run_initialized(wandb: mock.MagicMock) -> None:
+
+    wandb_kwargs = {
+        "project": "optuna",
+        "group": "summary",
+        "job_type": "logging",
+        "mode": "offline",
+    }
+
+    WeightsAndBiasesCallback(
+        metric_name="mse",
+        wandb_kwargs=wandb_kwargs,
+    )
+
+    wandb.init.assert_called_once_with(
+        project="optuna",
+        group="summary",
+        job_type="logging",
+        mode="offline",
+    )
+
+
+@mock.patch("optuna.integration.wandb.wandb")
+def test_attributes_set_on_epoch(wandb: mock.MagicMock) -> None:
+
+    wandb.config.update = mock.MagicMock()
+
+    wandbc = WeightsAndBiasesCallback()
+    study = optuna.create_study(direction="minimize")
+    study.optimize(_objective_func, n_trials=1, callbacks=[wandbc])
+
+    expected = {"direction": "MINIMIZE"}
+    wandb.config.update.assert_called_once_with(expected)
+
+
+@mock.patch("optuna.integration.wandb.wandb")
+def test_log_api_call_count(wandb: mock.Mock) -> None:
+
+    wandb.log = mock.MagicMock()
+
+    wandbc = WeightsAndBiasesCallback()
+    target_n_trials = 10
+    study = optuna.create_study()
+    study.optimize(_objective_func, n_trials=target_n_trials, callbacks=[wandbc])
+    assert wandb.log.call_count == target_n_trials
+
+
+@mock.patch("optuna.integration.wandb.wandb")
+def test_values_registered_on_epoch(wandb: mock.Mock) -> None:
+
+    wandb.log = mock.MagicMock()
+
+    wandbc = WeightsAndBiasesCallback()
+    study = optuna.create_study()
+    study.optimize(_objective_func, n_trials=1, callbacks=[wandbc])
+
+    kall = wandb.log.call_args
+    assert list(kall[0][0].keys()) == ["x", "y", "value"]
+    assert kall[1] == {"step": 0}