Abstract log_figure API that implements reasonable default for all lo…

CHANGELOG.md

@@ -9,6 +9,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Added
 
+- Added unified API for figure logging `log_figure` in Tensorboard, Comet, ML Flow, Neptune, Weights and Biases.
 
 ### Changed
 

docs/source/extensions/logging.rst

@@ -315,6 +315,20 @@ in the `hparams tab <https://pytorch.org/docs/stable/tensorboard.html#torch.util
 
 ----------
 
+***************
+Logging Figures
+***************
+
+When training a model, often it is very indicative to log figures, e.g. of the in- and output.
+For standard ``matplotlib.pyplot`` figures, Lightning has a unified API that works with Tensorboard, Comet, ML Flow, Neptune, Weights and Biases.
+
+.. code-block:: python
+
+    f = plt.figure()
+    logger.log_figure(name='dummy_figure', figure=f, step=0, close=True)
+
+----------
+
 *************
 Snapshot code
 *************

pytorch_lightning/loggers/base.py

@@ -25,7 +25,15 @@
 import torch
 
 from pytorch_lightning.core.lightning import LightningModule
-from pytorch_lightning.utilities import rank_zero_only
+from pytorch_lightning.utilities import _module_available, rank_zero_only
+
+_MATPLOTLIB_AVAILABLE = _module_available("matplotlib")
+
+if _MATPLOTLIB_AVAILABLE:
+    import matplotlib.pyplot as plt
+else:
+    from pytorch_lightning.utilities.mock_types import matplotlib as _matplotlib
+    plt = _matplotlib.pyplot
 
 
 def rank_zero_experiment(fn: Callable) -> Callable:
@@ -173,6 +181,23 @@ def log_metrics(self, metrics: Dict[str, float], step: Optional[int] = None):
         """
         pass
 
+    def log_figure(self, name: str, figure: plt.figure, step: Optional[int] = None, close: bool = True) -> None:
+        """
+        Logs a matplotlib figure.
+
+        Args:
+            name: name of the figure
+            figure: plt figure handle
+            step: step number at which the figure should be recorded
+            close: close figure after logging
+        """
+        import matplotlib.pyplot as plt
+
+        # Default is silent and not NotImplementedError because we want to support LoggerCollection
+        # where some loggers might others might not have implemented this method.
+        if close:
+            plt.close(figure)
+
     @staticmethod
     def _convert_params(params: Union[Dict[str, Any], Namespace]) -> Dict[str, Any]:
         # in case converting from namespace
@@ -377,6 +402,16 @@ def log_metrics(self, metrics: Dict[str, float], step: Optional[int] = None) ->
         for logger in self._logger_iterable:
             logger.log_metrics(metrics, step)
 
+    def log_figure(self, name: str, figure, step: Optional[int] = None, close: bool = True) -> None:
+        import matplotlib.pyplot as plt
+
+        for logger in self._logger_iterable:
+            # don't close in the individual loggers, but once at the end
+            logger.log_figure(name, figure, step=step, close=False)
+
+        if close:
+            plt.close(figure)
+
     def log_hyperparams(self, params: Union[Dict[str, Any], Namespace]) -> None:
         for logger in self._logger_iterable:
             logger.log_hyperparams(params)

pytorch_lightning/loggers/comet.py

@@ -49,6 +49,14 @@
     CometExperiment, CometExistingExperiment, CometOfflineExperiment = None, None, None
     API = None
 
+_MATPLOTLIB_AVAILABLE = _module_available("matplotlib")
+
+if _MATPLOTLIB_AVAILABLE:
+    import matplotlib.pyplot as plt
+else:
+    from pytorch_lightning.utilities.mock_types import matplotlib as _matplotlib
+    plt = _matplotlib.pyplot
+
 
 class CometLogger(LightningLoggerBase):
     r"""
@@ -252,6 +260,13 @@ def log_metrics(self, metrics: Dict[str, Union[torch.Tensor, float]], step: Opti
         metrics_without_epoch = self._add_prefix(metrics_without_epoch)
         self.experiment.log_metrics(metrics_without_epoch, step=step, epoch=epoch)
 
+    @rank_zero_only
+    def log_figure(self, name: str, figure: plt.figure, step: Optional[int] = None, close: bool = True) -> None:
+
+        self.experiment.log_figure(figure_name=name, figure=figure, step=step)
+        if close:
+            plt.close(figure)
+
     def reset_experiment(self):
         self._experiment = None
 

pytorch_lightning/loggers/mlflow.py

@@ -17,7 +17,9 @@
 """
 import logging
 import re
+import tempfile
 from argparse import Namespace
+from pathlib import Path
 from time import time
 from typing import Any, Dict, Optional, Union
 
@@ -35,6 +37,15 @@
     _MLFLOW_AVAILABLE = False
     mlflow, MlflowClient, context = None, None, None
 
+_MATPLOTLIB_AVAILABLE = _module_available("matplotlib")
+
+if _MATPLOTLIB_AVAILABLE:
+    import matplotlib.pyplot as plt
+else:
+    from pytorch_lightning.utilities.mock_types import matplotlib as _matplotlib
+    plt = _matplotlib.pyplot
+
+
 # before v1.1.0
 if hasattr(context, 'resolve_tags'):
     from mlflow.tracking.context import resolve_tags
@@ -93,6 +104,7 @@ def any_lightning_module_function_or_hook(self):
         prefix: A string to put at the beginning of metric keys.
         artifact_location: The location to store run artifacts. If not provided, the server picks an appropriate
             default.
+        figure_file_extension: File extension with which matplotlib saves figure
 
     Raises:
         ImportError:
@@ -109,6 +121,7 @@ def __init__(
         save_dir: Optional[str] = './mlruns',
         prefix: str = '',
         artifact_location: Optional[str] = None,
+        figure_file_extension='.png',
     ):
         if mlflow is None:
             raise ImportError(
@@ -128,6 +141,7 @@ def __init__(
         self._artifact_location = artifact_location
 
         self._mlflow_client = MlflowClient(tracking_uri)
+        self._figure_file_extension = figure_file_extension
 
     @property
     @rank_zero_experiment
@@ -204,6 +218,20 @@ def log_metrics(self, metrics: Dict[str, float], step: Optional[int] = None) ->
 
             self.experiment.log_metric(self.run_id, k, v, timestamp_ms, step)
 
+    @rank_zero_only
+    def log_figure(self, name: str, figure: plt.figure, step: Optional[int] = None, close: bool = True) -> None:
+
+        with tempfile.NamedTemporaryFile(suffix=self._figure_file_extension) as tmp_file:
+            figure.savefig(tmp_file)
+            self.experiment.log_artifact(
+                self.run_id,
+                tmp_file.name,
+                artifact_path=Path(self.save_dir) / ("figure_" + name)
+            )
+
+        if close:
+            plt.close(figure)
+
     @rank_zero_only
     def finalize(self, status: str = 'FINISHED') -> None:
         super().finalize(status)

pytorch_lightning/loggers/neptune.py

@@ -35,6 +35,14 @@
     # needed for test mocks, these tests shall be updated
     neptune, Experiment = None, None
 
+_MATPLOTLIB_AVAILABLE = _module_available("matplotlib")
+
+if _MATPLOTLIB_AVAILABLE:
+    import matplotlib.pyplot as plt
+else:
+    from pytorch_lightning.utilities.mock_types import matplotlib as _matplotlib
+    plt = _matplotlib.pyplot
+
 
 class NeptuneLogger(LightningLoggerBase):
     r"""
@@ -263,6 +271,15 @@ def log_metrics(self, metrics: Dict[str, Union[torch.Tensor, float]], step: Opti
             # Lighting does not always guarantee.
             self.log_metric(key, val)
 
+    @rank_zero_only
+    def log_figure(self, name: str, figure: plt.figure, step: Optional[int] = None, close: bool = True) -> None:
+        import matplotlib.pyplot as plt
+
+        description = f"step_{step}" if step is not None else None
+        self.experiment.log_image(name, figure, description=description)
+        if close:
+            plt.close(figure)
+
     @rank_zero_only
     def finalize(self, status: str) -> None:
         super().finalize(status)

pytorch_lightning/loggers/tensorboard.py

@@ -28,14 +28,23 @@
 from pytorch_lightning.core.lightning import LightningModule
 from pytorch_lightning.core.saving import save_hparams_to_yaml
 from pytorch_lightning.loggers.base import LightningLoggerBase, rank_zero_experiment
-from pytorch_lightning.utilities import _OMEGACONF_AVAILABLE, rank_zero_only, rank_zero_warn
+from pytorch_lightning.utilities import \
+    _OMEGACONF_AVAILABLE, _module_available, rank_zero_only, rank_zero_warn
 from pytorch_lightning.utilities.cloud_io import get_filesystem
 
 log = logging.getLogger(__name__)
 
 if _OMEGACONF_AVAILABLE:
     from omegaconf import Container, OmegaConf
 
+_MATPLOTLIB_AVAILABLE = _module_available("matplotlib")
+
+if _MATPLOTLIB_AVAILABLE:
+    import matplotlib.pyplot as plt
+else:
+    from pytorch_lightning.utilities.mock_types import matplotlib as _matplotlib
+    plt = _matplotlib.pyplot
+
 
 class TensorBoardLogger(LightningLoggerBase):
     r"""
@@ -209,6 +218,10 @@ def log_metrics(self, metrics: Dict[str, float], step: Optional[int] = None) ->
                     m = f'\n you tried to log {v} which is not currently supported. Try a dict or a scalar/tensor.'
                     raise ValueError(m) from ex
 
+    @rank_zero_only
+    def log_figure(self, name: str, figure: plt.figure, step: Optional[int] = None, close: bool = True) -> None:
+        self.experiment.add_figure(tag=name, figure=figure, global_step=step, close=close)
+
     @rank_zero_only
     def log_graph(self, model: LightningModule, input_array=None):
         if self._log_graph:

pytorch_lightning/loggers/wandb.py

@@ -37,6 +37,14 @@
     # needed for test mocks, these tests shall be updated
     wandb, Run = None, None
 
+_MATPLOTLIB_AVAILABLE = _module_available("matplotlib")
+
+if _MATPLOTLIB_AVAILABLE:
+    import matplotlib.pyplot as plt
+else:
+    from pytorch_lightning.utilities.mock_types import matplotlib as _matplotlib
+    plt = _matplotlib.pyplot
+
 
 class WandbLogger(LightningLoggerBase):
     r"""
@@ -199,6 +207,14 @@ def log_metrics(self, metrics: Dict[str, float], step: Optional[int] = None) ->
         else:
             self.experiment.log(metrics)
 
+    @rank_zero_only
+    def log_figure(self, name: str, figure: plt.figure, step: Optional[int] = None, close: bool = True) -> None:
+        import matplotlib.pyplot as plt
+
+        self.experiment.log({name: wandb.Image(figure)}, step=step)
+        if close:
+            plt.close(figure)
+
     @property
     def save_dir(self) -> Optional[str]:
         return self._save_dir

pytorch_lightning/utilities/mock_types.py

@@ -0,0 +1,19 @@
+# Copyright The PyTorch Lightning team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+class matplotlib:
+    class pyplot:
+        figure = None
+        close = None

tests/helpers/plotting.py

@@ -0,0 +1,19 @@
+import numpy as np
+
+from pytorch_lightning.utilities import _module_available
+
+_MATPLOTLIB_AVAILABLE = _module_available("matplotlib")
+if _MATPLOTLIB_AVAILABLE:
+    import matplotlib.pyplot as plt
+else:
+    class plt:
+        figure = None
+
+
+def dummy_figure() -> plt.figure:
+    """Dummy figure to test logging capability of figures for loggers."""
+
+    f = plt.figure()
+    plt.plot(np.linspace(0., 1., 100), np.linspace(0., 10., 100) ** 2)
+
+    return f

tests/loggers/test_all.py

@@ -24,14 +24,16 @@
 from pytorch_lightning import Callback, Trainer
 from pytorch_lightning.loggers import (
     CometLogger,
+    CSVLogger,
     MLFlowLogger,
     NeptuneLogger,
     TensorBoardLogger,
     TestTubeLogger,
     WandbLogger,
 )
 from pytorch_lightning.loggers.base import DummyExperiment
-from tests.helpers import BoringModel
+from pytorch_lightning.utilities import _module_available
+from tests.helpers import BoringModel, plotting
 from tests.helpers.runif import RunIf
 from tests.loggers.test_comet import _patch_comet_atexit
 from tests.loggers.test_mlflow import mock_mlflow_run_creation
@@ -404,3 +406,40 @@ def test_logger_with_prefix_all(tmpdir, monkeypatch):
         wandb.init().step = 0
         logger.log_metrics({"test": 1.0}, step=0)
         logger.experiment.log.assert_called_once_with({'tmp-test': 1.0, 'trainer/global_step': 0})
+
+
+@pytest.mark.skipif(
+    not _module_available("matplotlib"),
+    reason="close figure test requires matplotlib to be installed.")
+@pytest.mark.parametrize("close", [True, False])
+@pytest.mark.parametrize("logger_class", [
+    CometLogger,
+    CSVLogger,
+    MLFlowLogger,
+    NeptuneLogger,
+    TensorBoardLogger,
+    # Wandb has its own close_figure test
+])
+def test_logger_close_figure_all(tmpdir, monkeypatch, logger_class, close):
+    _patch_comet_atexit(monkeypatch)
+    try:
+        _test_logger_close_figure(tmpdir, monkeypatch, logger_class, close)
+    except (ImportError, ModuleNotFoundError):
+        pytest.xfail(f"pickle test requires {logger_class.__class__} dependencies to be installed.")
+
+
+def _test_logger_close_figure(tmpdir, monkeypatch, logger_class, close):
+    _patch_comet_atexit(monkeypatch)
+
+    logger_args = _get_logger_args(logger_class, tmpdir)
+    logger = logger_class(**logger_args)
+
+    f = plotting.dummy_figure()
+
+    with mock.patch('matplotlib.pyplot.close') as plt_close:
+        logger.log_figure('dummy', f, 0, close=close)
+
+    if close:
+        plt_close.assert_called_once_with(f)
+    else:
+        plt_close.assert_not_called()