Add support for WandbLogger (#1176)

* add WandbLogger * update logger * Update tests, docstring * add requirements-wandb.txt * Update catalyst/loggers/wandb.py Co-authored-by: Sergey Kolesnikov <scitator@gmail.com> * update wandblogger * PEP8 format * reformat * FIX test example * Update catalyst/tests/test_finetune2.py Co-authored-by: Sergey Kolesnikov <scitator@gmail.com> * Update hparams logging * fix typo * Allow config update, update test * Update metric names * Update metric name format * Update CHANGELOG.md * Update CHANGELOG.md * Update CHANGELOG.md * Update loggers.rst Co-authored-by: Sergey Kolesnikov <scitator@gmail.com>
catalyst-team · Apr 30, 2021 · 6cfa27f · 6cfa27f
1 parent f07a040
commit 6cfa27f
Show file tree

Hide file tree

Showing 7 changed files with 218 additions and 1 deletion.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -9,6 +9,8 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Added
 
+
+- Weights and Biases Logger (``WandbLogger``) ([#1176](https://github.com/catalyst-team/catalyst/pull/1176))
 - Neptune Logger (``NeptuneLogger``) ([#1196](https://github.com/catalyst-team/catalyst/pull/1196)
 - `log_artifact` method for logging arbitrary files like audio, video, or model weights to `ILogger` and `IRunner` ([#1196](https://github.com/catalyst-team/catalyst/pull/1196)
 

diff --git a/catalyst/loggers/__init__.py b/catalyst/loggers/__init__.py
@@ -9,6 +9,9 @@
 if SETTINGS.mlflow_required:
     from catalyst.loggers.mlflow import MLflowLogger
 
+if SETTINGS.wandb_required:
+    from catalyst.loggers.wandb import WandbLogger
+
 if SETTINGS.neptune_required:
     from catalyst.loggers.neptune import NeptuneLogger
 
@@ -18,5 +21,8 @@
 if SETTINGS.mlflow_required:
     __all__ += ["MLflowLogger"]
 
+if SETTINGS.wandb_required:
+    __all__ += ["WandbLogger"]
+
 if SETTINGS.neptune_required:
     __all__ += ["NeptuneLogger"]
diff --git a/catalyst/loggers/wandb.py b/catalyst/loggers/wandb.py
@@ -0,0 +1,180 @@
+from typing import Any, Dict, Optional
+
+import numpy as np
+
+from catalyst.core.logger import ILogger
+from catalyst.settings import SETTINGS
+
+if SETTINGS.wandb_required:
+    import wandb
+
+
+class WandbLogger(ILogger):
+    """Wandb logger for parameters, metrics, images and other artifacts.
+
+    W&B documentation: https://docs.wandb.com
+
+    Args:
+        Project: Name of the project in W&B to log to.
+        name: Name of the run in W&B to log to.
+        config: Configuration Dictionary for the experiment.
+        entity: Name of W&B entity(team) to log to.
+
+    Python API examples:
+
+    .. code-block:: python
+
+        from catalyst import dl
+
+        runner = dl.SupervisedRunner()
+        runner.train(
+            ...,
+            loggers={"wandb": dl.WandbLogger(project="wandb_test", name="expeirment_1")}
+        )
+
+    .. code-block:: python
+
+        from catalyst import dl
+
+        class CustomRunner(dl.IRunner):
+            # ...
+
+            def get_loggers(self):
+                return {
+                    "console": dl.ConsoleLogger(),
+                    "wandb": dl.WandbLogger(project="wandb_test", name="experiment_1")
+                }
+
+            # ...
+
+        runner = CustomRunner().run()
+
+    Config API example:
+
+    .. code-block:: yaml
+
+        loggers:
+            wandb:
+                _target_: WandbLogger
+                project: test_exp
+                name: test_run
+        ...
+
+    Hydra API example:
+
+    .. code-block:: yaml
+
+        loggers:
+            wandb:
+                _target_: catalyst.dl.WandbLogger
+                project: test_exp
+                name: test_run
+        ...
+    """
+
+    def __init__(
+        self, project: str, name: Optional[str] = None, entity: Optional[str] = None,
+    ) -> None:
+        self.project = project
+        self.name = name
+        self.entity = entity
+        self.run = wandb.init(
+            project=self.project, name=self.name, entity=self.entity, allow_val_change=True
+        )
+
+    def _log_metrics(self, metrics: Dict[str, float], step: int, loader_key: str, prefix=""):
+        for key, value in metrics.items():
+            self.run.log({f"{key}_{prefix}/{loader_key}": value}, step=step)
+
+    def log_metrics(
+        self,
+        metrics: Dict[str, Any],
+        scope: str = None,
+        # experiment info
+        run_key: str = None,
+        global_epoch_step: int = 0,
+        global_batch_step: int = 0,
+        global_sample_step: int = 0,
+        # stage info
+        stage_key: str = None,
+        stage_epoch_len: int = 0,
+        stage_epoch_step: int = 0,
+        stage_batch_step: int = 0,
+        stage_sample_step: int = 0,
+        # loader info
+        loader_key: str = None,
+        loader_batch_len: int = 0,
+        loader_sample_len: int = 0,
+        loader_batch_step: int = 0,
+        loader_sample_step: int = 0,
+    ) -> None:
+        """Logs batch and epoch metrics to wandb."""
+        if scope == "batch":
+            metrics = {k: float(v) for k, v in metrics.items()}
+            self._log_metrics(
+                metrics=metrics, step=global_epoch_step, loader_key=loader_key, prefix="batch"
+            )
+        elif scope == "loader":
+            self._log_metrics(
+                metrics=metrics, step=global_epoch_step, loader_key=loader_key, prefix="epoch",
+            )
+        elif scope == "epoch":
+            loader_key = "_epoch_"
+            per_loader_metrics = metrics[loader_key]
+            self._log_metrics(
+                metrics=per_loader_metrics,
+                step=global_epoch_step,
+                loader_key=loader_key,
+                prefix="epoch",
+            )
+
+    def log_image(
+        self,
+        tag: str,
+        image: np.ndarray,
+        scope: str = None,
+        # experiment info
+        run_key: str = None,
+        global_epoch_step: int = 0,
+        global_batch_step: int = 0,
+        global_sample_step: int = 0,
+        # stage info
+        stage_key: str = None,
+        stage_epoch_len: int = 0,
+        stage_epoch_step: int = 0,
+        stage_batch_step: int = 0,
+        stage_sample_step: int = 0,
+        # loader info
+        loader_key: str = None,
+        loader_batch_len: int = 0,
+        loader_sample_len: int = 0,
+        loader_batch_step: int = 0,
+        loader_sample_step: int = 0,
+    ) -> None:
+        """Logs image to the logger."""
+        self.run.log(
+            {f"{tag}_scope_{scope}_epoch_{global_epoch_step}.png": wandb.Image(image)},
+            step=global_epoch_step,
+        )
+
+    def log_hparams(
+        self,
+        hparams: Dict,
+        scope: str = None,
+        # experiment info
+        run_key: str = None,
+        stage_key: str = None,
+    ) -> None:
+        """Logs hyperparameters to the logger."""
+        self.run.config.update(hparams)
+
+    def flush_log(self) -> None:
+        """Flushes the logger."""
+        pass
+
+    def close_log(self) -> None:
+        """Closes the logger."""
+        self.run.finish()
+
+
+__all__ = ["WandbLogger"]
diff --git a/catalyst/settings.py b/catalyst/settings.py
@@ -131,6 +131,15 @@ def _is_mlflow_available():
         return False
 
 
+def _is_wandb_available():
+    try:
+        import wandb  # noqa: F401
+
+        return True
+    except ImportError:
+        return False
+
+
 def _is_neptune_available():
     try:
         import neptune.new as neptune  # noqa: F401
@@ -177,7 +186,7 @@ def __init__(  # noqa: D107
         # alchemy_required: Optional[bool] = None,
         neptune_required: Optional[bool] = None,
         mlflow_required: Optional[bool] = None,
-        # wandb_required: Optional[bool] = None,
+        wandb_required: Optional[bool] = None,
         # [extras]
         use_lz4: Optional[bool] = None,
         use_pyarrow: Optional[bool] = None,
@@ -268,6 +277,13 @@ def __init__(  # noqa: D107
             "catalyst[mlflow] is not available, to install it, "
             "run `pip install catalyst[mlflow]`.",
         )
+
+        self.wandb_required: bool = _get_optional_value(
+            wandb_required,
+            _is_wandb_available,
+            "wandb is not available, to install it, " "run `pip install wandb`.",
+        )
+
         # self.wandb_required: bool = wandb_required
 
         # [extras]

diff --git a/catalyst/tests/test_finetune2.py b/catalyst/tests/test_finetune2.py
@@ -32,12 +32,17 @@ def get_loggers(self):
         }
         if SETTINGS.mlflow_required:
             loggers["mlflow"] = dl.MLflowLogger(experiment=self._name)
+
+        if SETTINGS.wandb_required:
+            loggers["wandb"] = dl.WandbLogger(project="catalyst_test", name=self._name)
+
         if SETTINGS.neptune_required:
             loggers["neptune"] = dl.NeptuneLogger(
                 base_namespace="catalyst-tests",
                 api_token="ANONYMOUS",
                 project="common/catalyst-integration",
             )
+
         return loggers
 
     @property

diff --git a/docs/api/loggers.rst b/docs/api/loggers.rst
@@ -36,6 +36,13 @@ TensorboardLogger
     :undoc-members:
     :show-inheritance:
 
+WandbLogger
+--------------------------------
+.. autoclass:: catalyst.loggers.wandb.WandbLogger
+    :exclude-members: __init__
+    :undoc-members:
+    :show-inheritance:
+
 NeptuneLogger
 --------------------------------
 .. autoclass:: catalyst.loggers.neptune.NeptuneLogger

diff --git a/requirements/requirements-wandb.txt b/requirements/requirements-wandb.txt
@@ -0,0 +1 @@
+wandb