From 6cfa27f8815daa121ef77232e5f4c0103ad13d23 Mon Sep 17 00:00:00 2001 From: Ayush Chaurasia Date: Fri, 30 Apr 2021 20:33:35 +0530 Subject: [PATCH] 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 * update wandblogger * PEP8 format * reformat * FIX test example * Update catalyst/tests/test_finetune2.py Co-authored-by: Sergey Kolesnikov * 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 --- CHANGELOG.md | 2 + catalyst/loggers/__init__.py | 6 + catalyst/loggers/wandb.py | 180 ++++++++++++++++++++++++++++ catalyst/settings.py | 18 ++- catalyst/tests/test_finetune2.py | 5 + docs/api/loggers.rst | 7 ++ requirements/requirements-wandb.txt | 1 + 7 files changed, 218 insertions(+), 1 deletion(-) create mode 100644 catalyst/loggers/wandb.py create mode 100644 requirements/requirements-wandb.txt diff --git a/CHANGELOG.md b/CHANGELOG.md index ade6c3885d..bbe6ca93ec 100644 --- 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 index 7d93ad5ba5..85f02c8a34 100644 --- 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 new file mode 100644 index 0000000000..edbd3978e8 --- /dev/null +++ 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 index c6aaedf6c0..98ba1207cc 100644 --- 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 index 74182f8ffa..83654e0426 100644 --- 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 index f2a0946bbe..893c30ab1f 100644 --- 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 new file mode 100644 index 0000000000..8a70eb6fdb --- /dev/null +++ b/requirements/requirements-wandb.txt @@ -0,0 +1 @@ +wandb \ No newline at end of file