Neptune Logger Improvements (Lightning-AI#1084)

* removed project and experiment from getstate * added tests for closing experiment, updated token in example to user neptuner * updated teoken * Update neptune.py added a link to example experiment * added exmaple experiment link * dropped duplication * flake fixes * merged with master, added changes information to CHANGELOG
lezwon · Mar 14, 2020 · 3ad6169 · 3ad6169
1 parent 2232eb3
commit 3ad6169
Show file tree

Hide file tree

Showing 3 changed files with 123 additions and 34 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -68,6 +68,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Changed
 
+- Improved `NeptuneLogger` by adding `close_after_fit` argument to allow logging after training([#908](https://github.com/PyTorchLightning/pytorch-lightning/pull/1084))
 - Changed default TQDM to use `tqdm.auto` for prettier outputs in IPython notebooks ([#752](https://github.com/PyTorchLightning/pytorch-lightning/pull/752))
 - Changed `pytorch_lightning.logging` to `pytorch_lightning.loggers` ([#767](https://github.com/PyTorchLightning/pytorch-lightning/pull/767))
 - Moved the default `tqdm_dict` definition from Trainer to `LightningModule`, so it can be overridden by the user ([#749](https://github.com/PyTorchLightning/pytorch-lightning/pull/749))

diff --git a/pytorch_lightning/loggers/neptune.py b/pytorch_lightning/loggers/neptune.py
@@ -1,5 +1,5 @@
 """
-Log using `neptune-logger <https://www.neptune.ml>`_
+Log using `neptune-logger <https://neptune.ai>`_
 
 .. _neptune:
 
@@ -30,12 +30,13 @@ class NeptuneLogger(LightningLoggerBase):
     """
 
     def __init__(self, api_key: Optional[str] = None, project_name: Optional[str] = None,
-                 offline_mode: bool = False, experiment_name: Optional[str] = None,
+                 close_after_fit: Optional[bool] = True, offline_mode: bool = False,
+                 experiment_name: Optional[str] = None,
                  upload_source_files: Optional[List[str]] = None, params: Optional[Dict[str, Any]] = None,
                  properties: Optional[Dict[str, Any]] = None, tags: Optional[List[str]] = None, **kwargs):
         r"""
 
-        Initialize a neptune.ml logger.
+        Initialize a neptune.ai logger.
 
         .. note:: Requires either an API Key (online mode) or a local directory path (offline mode)
 
@@ -44,10 +45,11 @@ def __init__(self, api_key: Optional[str] = None, project_name: Optional[str] =
             # ONLINE MODE
             from pytorch_lightning.loggers import NeptuneLogger
             # arguments made to NeptuneLogger are passed on to the neptune.experiments.Experiment class
+            # We are using an api_key for the anonymous user "neptuner" but you can use your own.
 
             neptune_logger = NeptuneLogger(
-                api_key=os.environ["NEPTUNE_API_TOKEN"],
-                project_name="USER_NAME/PROJECT_NAME",
+                api_key="ANONYMOUS"
+                project_name="shared/pytorch-lightning-integration",
                 experiment_name="default", # Optional,
                 params={"max_epochs": 10}, # Optional,
                 tags=["pytorch-lightning","mlp"] # Optional,
@@ -85,40 +87,91 @@ def any_lightning_module_function_or_hook(...):
                 self.logger.experiment.log_artifact("model_checkpoint.pt", prediction_image) # log model checkpoint
                 self.logger.experiment.whatever_neptune_supports(...)
 
+        If you want to log objects after the training is finished use close_after_train=False:
+
+        .. code-block:: python
+
+            neptune_logger = NeptuneLogger(
+                ...
+                close_after_fit=False,
+                ...)
+            trainer = Trainer(logger=neptune_logger)
+            trainer.fit()
+
+            # Log test metrics
+            trainer.test(model)
+
+            # Log additional metrics
+            from sklearn.metrics import accuracy_score
+
+            accuracy = accuracy_score(y_true, y_pred)
+            neptune_logger.experiment.log_metric('test_accuracy', accuracy)
+
+            # Log charts
+            from scikitplot.metrics import plot_confusion_matrix
+            import matplotlib.pyplot as plt
+
+            fig, ax = plt.subplots(figsize=(16, 12))
+            plot_confusion_matrix(y_true, y_pred, ax=ax)
+            neptune_logger.experiment.log_image('confusion_matrix', fig)
+
+            # Save checkpoints folder
+            neptune_logger.experiment.log_artifact('my/checkpoints')
+
+            # When you are done, stop the experiment
+            neptune_logger.experiment.stop()
+
+        You can go and see an example experiment here:
+        https://ui.neptune.ai/o/shared/org/pytorch-lightning-integration/e/PYTOR-66/charts
+
         Args:
-            api_key (str | None): Required in online mode. Neputne API token, found on https://neptune.ml.
+            api_key: Required in online mode.
+                Neputne API token, found on https://neptune.ai
                 Read how to get your API key
-                https://docs.neptune.ml/python-api/tutorials/get-started.html#copy-api-token.
-            project_name (str): Required in online mode. Qualified name of a project in a form of
+                https://docs.neptune.ai/python-api/tutorials/get-started.html#copy-api-token.
+                It is recommended to keep it in the `NEPTUNE_API_TOKEN`
+                environment variable and then you can leave `api_key=None`
+            project_name: Required in online mode. Qualified name of a project in a form of
                "namespace/project_name" for example "tom/minst-classification".
                If None, the value of NEPTUNE_PROJECT environment variable will be taken.
-               You need to create the project in https://neptune.ml first.
-            offline_mode (bool): Optional default False. If offline_mode=True no logs will be send to neptune.
-               Usually used for debug purposes.
-            experiment_name (str|None): Optional. Editable name of the experiment.
-               Name is displayed in the experiment’s Details (Metadata section) and in experiments view as a column.
-            upload_source_files (list|None): Optional. List of source files to be uploaded.
-               Must be list of str or single str. Uploaded sources are displayed in the experiment’s Source code tab.
+               You need to create the project in https://neptune.ai first.
+            offline_mode: Optional default False. If offline_mode=True no logs will be send
+               to neptune. Usually used for debug purposes.
+            close_after_fit: Optional default True. If close_after_fit=False the experiment
+               will not be closed after training and additional metrics,
+               images or artifacts can be logged. Also, remember to close the experiment explicitly
+               by running neptune_logger.experiment.stop().
+            experiment_name: Optional. Editable name of the experiment.
+               Name is displayed in the experiment’s Details (Metadata section) and i
+               n experiments view as a column.
+            upload_source_files: Optional. List of source files to be uploaded.
+               Must be list of str or single str. Uploaded sources are displayed
+               in the experiment’s Source code tab.
                If None is passed, Python file from which experiment was created will be uploaded.
-               Pass empty list ([]) to upload no files. Unix style pathname pattern expansion is supported.
+               Pass empty list ([]) to upload no files.
+               Unix style pathname pattern expansion is supported.
                For example, you can pass '\*.py'
                 to upload all python source files from the current directory.
                For recursion lookup use '\**/\*.py' (for Python 3.5 and later).
                For more information see glob library.
-            params (dict|None): Optional. Parameters of the experiment. After experiment creation params are read-only.
-               Parameters are displayed in the experiment’s Parameters section and each key-value pair can be
-               viewed in experiments view as a column.
-            properties (dict|None): Optional default is {}. Properties of the experiment.
-               They are editable after experiment is created. Properties are displayed in the experiment’s Details and
+            params: Optional. Parameters of the experiment.
+               After experiment creation params are read-only.
+               Parameters are displayed in the experiment’s Parameters section and
+               each key-value pair can be viewed in experiments view as a column.
+            properties: Optional default is {}. Properties of the experiment.
+               They are editable after experiment is created.
+               Properties are displayed in the experiment’s Details and
                each key-value pair can be viewed in experiments view as a column.
-            tags (list|None): Optional default []. Must be list of str. Tags of the experiment.
+            tags: Optional default []. Must be list of str. Tags of the experiment.
                They are editable after experiment is created (see: append_tag() and remove_tag()).
-               Tags are displayed in the experiment’s Details and can be viewed in experiments view as a column.
+               Tags are displayed in the experiment’s Details and can be viewed
+               in experiments view as a column.
         """
         super().__init__()
         self.api_key = api_key
         self.project_name = project_name
         self.offline_mode = offline_mode
+        self.close_after_fit = close_after_fit
         self.experiment_name = experiment_name
         self.upload_source_files = upload_source_files
         self.params = params
@@ -138,6 +191,12 @@ def any_lightning_module_function_or_hook(...):
 
         log.info(f'NeptuneLogger was initialized in {self.mode} mode')
 
+    def __getstate__(self):
+        state = self.__dict__.copy()
+        # cannot be pickled
+        state['_experiment'] = None
+        return state
+
     @property
     def experiment(self) -> Experiment:
         r"""
@@ -150,15 +209,14 @@ def experiment(self) -> Experiment:
 
         """
 
-        if self._experiment is not None:
-            return self._experiment
-        else:
-            self._experiment = neptune.create_experiment(name=self.experiment_name,
-                                                         params=self.params,
-                                                         properties=self.properties,
-                                                         tags=self.tags,
-                                                         upload_source_files=self.upload_source_files,
-                                                         **self._kwargs)
+        if self._experiment is None:
+            self._experiment = neptune.create_experiment(
+                name=self.experiment_name,
+                params=self.params,
+                properties=self.properties,
+                tags=self.tags,
+                upload_source_files=self.upload_source_files,
+                **self._kwargs)
         return self._experiment
 
     @rank_zero_only
@@ -184,7 +242,8 @@ def log_metrics(
 
     @rank_zero_only
     def finalize(self, status: str) -> None:
-        self.experiment.stop()
+        if self.close_after_fit:
+            self.experiment.stop()
 
     @property
     def name(self) -> str:

diff --git a/tests/loggers/test_neptune.py b/tests/loggers/test_neptune.py
@@ -1,5 +1,6 @@
 import pickle
-from unittest.mock import patch
+
+from unittest.mock import patch, MagicMock
 
 import torch
 
@@ -96,3 +97,31 @@ def test_neptune_pickle(tmpdir):
     pkl_bytes = pickle.dumps(trainer)
     trainer2 = pickle.loads(pkl_bytes)
     trainer2.logger.log_metrics({'acc': 1.0})
+
+
+def test_neptune_leave_open_experiment_after_fit(tmpdir):
+    """Verify that neptune experiment was closed after training"""
+    tutils.reset_seed()
+
+    hparams = tutils.get_hparams()
+    model = LightningTestModel(hparams)
+
+    def _run_training(logger):
+        logger._experiment = MagicMock()
+
+        trainer_options = dict(
+            default_save_path=tmpdir,
+            max_epochs=1,
+            train_percent_check=0.05,
+            logger=logger
+        )
+        trainer = Trainer(**trainer_options)
+        trainer.fit(model)
+        return logger
+
+    logger_close_after_fit = _run_training(NeptuneLogger(offline_mode=True))
+    assert logger_close_after_fit._experiment.stop.call_count == 1
+
+    logger_open_after_fit = _run_training(
+        NeptuneLogger(offline_mode=True, close_after_fit=False))
+    assert logger_open_after_fit._experiment.stop.call_count == 0