Repository: Lightning-AI/pytorch-lightning

Issue:
NeptuneLogger doesn't work with distributed_backend='ddp'
<!-- 
### Common bugs:
1. Tensorboard not showing in Jupyter-notebook see [issue 79](https://github.com/PyTorchLightning/pytorch-lightning/issues/79).    
2. PyTorch 1.1.0 vs 1.2.0 support [see FAQ](https://github.com/PyTorchLightning/pytorch-lightning#faq)    
-->

## 🐛 Bug

<!-- A clear and concise description of what the bug is. -->
When using NeptuneLogger with distributed_backend='ddp' and running it on a single node with two GPUs, I find an error like this.
```
Traceback (most recent call last):
  File "pl.py", line 146, in <module>
    main()
  File "pl.py", line 103, in main
    trainer.fit(model)
  File "/home/hirune/anaconda3/envs/PANDA/lib/python3.7/site-packages/pytorch_lightning/trainer/trainer.py", line 751, in fit
    mp.spawn(self.ddp_train, nprocs=self.num_processes, args=(model,))
  File "/home/hirune/anaconda3/envs/PANDA/lib/python3.7/site-packages/torch/multiprocessing/spawn.py", line 200, in spawn
    return start_processes(fn, args, nprocs, join, daemon, start_method='spawn')
  File "/home/hirune/anaconda3/envs/PANDA/lib/python3.7/site-packages/torch/multiprocessing/spawn.py", line 158, in start_processes
    while not context.join():
  File "/home/hirune/anaconda3/envs/PANDA/lib/python3.7/site-packages/torch/multiprocessing/spawn.py", line 119, in join
    raise Exception(msg)
Exception:

-- Process 1 terminated with the following error:
Traceback (most recent call last):
  File "/home/hirune/anaconda3/envs/PANDA/lib/python3.7/site-packages/torch/multiprocessing/spawn.py", line 20, in _wrap
    fn(i, *args)
  File "/home/hirune/anaconda3/envs/PANDA/lib/python3.7/site-packages/pytorch_lightning/trainer/distrib_data_parallel.py", line 370, in ddp_train
    self.run_pretrain_routine(model)
  File "/home/hirune/anaconda3/envs/PANDA/lib/python3.7/site-packages/pytorch_lightning/trainer/trainer.py", line 871, in run_pretrain_routine
    self.configure_checkpoint_callback()
  File "/home/hirune/anaconda3/envs/PANDA/lib/python3.7/site-packages/pytorch_lightning/trainer/callback_config.py", line 54, in configure_checkpoint_callback
    self.logger.name,
  File "/home/hirune/anaconda3/envs/PANDA/lib/python3.7/site-packages/pytorch_lightning/loggers/neptune.py", line 267, in name
    return self.experiment.name
  File "/home/hirune/anaconda3/envs/PANDA/lib/python3.7/site-packages/pytorch_lightning/loggers/neptune.py", line 230, in experiment
    **self._kwargs)
  File "/home/hirune/anaconda3/envs/PANDA/lib/python3.7/site-packages/neptune/__init__.py", line 222, in create_experiment
    raise Uninitialized()
neptune.exceptions.Uninitialized: You must initialize neptune-client first. For more information, please visit: https://github.com/neptune-ai/neptune-client#initialize-neptune
```
And I found a similar error with CommetLogger

### To Reproduce

Steps to reproduce the behavior:

Run the following code on a machine with two GPUs.
This code is a slightly modified version of what was on this page.
https://docs.neptune.ai/integrations/pytorch_lightning.html

<!-- If you have a code sample, error messages, stack traces, please provide it here as well -->


#### Code sample
<!-- Ideally attach a minimal code sample to reproduce the decried issue. 
Minimal means having the shortest code but still preserving the bug. -->
```
import os

import torch
from torch.nn import functional as F
from torch.utils.data import DataLoader
from torchvision.datasets import MNIST
from torchvision import transforms

import pytorch_lightning as pl

MAX_EPOCHS=20
LR=0.01
BATCHSIZE=32
CHECKPOINTS_DIR = 'my_models/checkpoints/7'

class CoolSystem(pl.LightningModule):

    def __init__(self):
        super(CoolSystem, self).__init__()
        # not the best model...
        self.l1 = torch.nn.Linear(28 * 28, 10)

    def forward(self, x):
        return torch.relu(self.l1(x.view(x.size(0), -1)))

    def training_step(self, batch, batch_idx):
        # REQUIRED
        x, y = batch
        y_hat = self.forward(x)
        loss = F.cross_entropy(y_hat, y)
        tensorboard_logs = {'train_loss': loss}
        return {'loss': loss, 'log': tensorboard_logs}

    def validation_step(self, batch, batch_idx):
        # OPTIONAL
        x, y = batch
        y_hat = self.forward(x)
        return {'val_loss': F.cross_entropy(y_hat, y)}

    def validation_end(self, outputs):
        # OPTIONAL
        avg_loss = torch.stack([x['val_loss'] for x in outputs]).mean()
        tensorboard_logs = {'val_loss': avg_loss}
        return {'avg_val_loss': avg_loss, 'log': tensorboard_logs}

    def test_step(self, batch, batch_idx):
        # OPTIONAL
        x, y = batch
        y_hat = self.forward(x)
        return {'test_loss': F.cross_entropy(y_hat, y)}

    def test_end(self, outputs):
        # OPTIONAL
        avg_loss = torch.stack([x['test_loss'] for x in outputs]).mean()
        tensorboard_logs = {'test_loss': avg_loss}
        return {'avg_test_loss': avg_loss, 'log': tensorboard_logs}

    def configure_optimizers(self):
        # REQUIRED
        # can return multiple optimizers and learning_rate schedulers
        # (LBFGS it is automatically supported, no need for closure function)
        return torch.optim.Adam(self.parameters(), lr=LR)

    @pl.data_loader
    def train_dataloader(self):
        # REQUIRED
        return DataLoader(MNIST(os.getcwd(), train=True, download=True, transform=transforms.ToTensor()), batch_size=BATCHSIZE)

    @pl.data_loader
    def val_dataloader(self):
        # OPTIONAL
        return DataLoader(MNIST(os.getcwd(), train=True, download=True, transform=transforms.ToTensor()), batch_size=BATCHSIZE)

    @pl.data_loader
    def test_dataloader(self):
        # OPTIONAL
        return DataLoader(MNIST(os.getcwd(), train=False, download=True, transform=transforms.ToTensor()), batch_size=BATCHSIZE)


from pytorch_lightning.loggers.neptune import NeptuneLogger
def main():
    neptune_logger = NeptuneLogger(
        api_key="ANONYMOUS",
        project_name="shared/pytorch-lightning-integration",
        close_after_fit=False,
        experiment_name="default",  # Optional,
        params={"max_epochs": MAX_EPOCHS,
                "batch_size": BATCHSIZE,
                "lr": LR}, # Optional,
        tags=["pytorch-lightning", "mlp"]  # Optional,
    )
    model_checkpoint = pl.callbacks.ModelCheckpoint(filepath=CHECKPOINTS_DIR)
    
    from pytorch_lightning import Trainer
    
    model = CoolSystem()
    trainer = Trainer(max_epochs=MAX_EPOCHS,
                      logger=neptune_logger,
                      checkpoint_callback=model_checkpoint,
                      gpus=-1,
                      distributed_backend='ddp',
                      )
    trainer.fit(model)
    trainer.test(model)
    
    # Get predictions on external test
    import numpy as np
    
    model.freeze()
    test_loader = DataLoader(MNIST(os.getcwd(), train=False, download=True, transform=transforms.ToTensor()), batch_size=256)
    
    y_true, y_pred = [],[]
    for i, (x, y) in enumerate(test_loader):
        y_hat = model.forward(x).argmax(axis=1).cpu().detach().numpy()
        y = y.cpu().detach().numpy()
    
        y_true.append(y)
        y_pred.append(y_hat)
    
        if i == len(test_loader):
            break
    y_true = np.hstack(y_true)
    y_pred = np.hstack(y_pred)
    
    # 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(CHECKPOINTS_DIR)
    
    # You can stop the experiment
    neptune_logger.experiment.stop()

if __name__ == "__main__":
        main()

```
### Expected behavior

<!-- A clear and concise description of what you expected to happen. -->

### Environment
* CUDA:
        - GPU:
                - GeForce GTX TITAN X
                - GeForce GTX TITAN X
        - available:         True
        - version:           10.1
* Packages:
        - numpy:             1.18.1
        - pyTorch_debug:     False
        - pyTorch_version:   1.5.0
        - pytorch-lightning: 0.7.5
        - tensorboard:       2.2.1
        - tqdm:              4.42.1
* System:
        - OS:                Linux
        - architecture:
                - 64bit
                -
        - processor:         x86_64
        - python:            3.7.6
        - version:           #86-Ubuntu SMP Fri Jan 17 17:24:28 UTC 2020

### Additional context

<!-- Add any other context about the problem here. -->



Discussion:
Hi! thanks for your contribution!, great first issue!
@jakubczakon pls ^^
Hi @hirune924 and thanks for raising it!

I've already notified the dev team and @pitercl will get back to you once we have a solution.
@jakubczakon I assume that this is also Neptune issue only, not related to PL, right?
@hirune924 mentioned:

> And I found a similar error with CommetLogger

So I think it may be a more general problem @Borda, 
Hi @hirune924, @Borda, just a quick update on this: from my initial tests, our Neptune logger and multiprocessing don't mix well. Tomorrow I'll dig a bit deeper into that and see if/how it can be remedied.

@Borda, I don't know yet if this is more on Neptune side or PL side.
Another quick update: I have a solution to this and will create a PR with a fix today or tomorrow.

Resolving patch:
diff --git a/pytorch_lightning/loggers/neptune.py b/pytorch_lightning/loggers/neptune.py
--- a/pytorch_lightning/loggers/neptune.py
+++ b/pytorch_lightning/loggers/neptune.py
@@ -32,7 +32,7 @@ class NeptuneLogger(LightningLoggerBase):
 
     The Neptune logger can be used in the online mode or offline (silent) mode.
     To log experiment data in online mode, :class:`NeptuneLogger` requires an API key.
-    In offline mode, Neptune will log to a local directory.
+    In offline mode, the logger does not connect to Neptune.
 
     **ONLINE MODE**
 
@@ -83,7 +83,7 @@ class NeptuneLogger(LightningLoggerBase):
     ...         self.logger.experiment.log_artifact('model_checkpoint.pt', ...)
     ...         self.logger.experiment.whatever_neptune_supports(...)
 
-    If you want to log objects after the training is finished use ``close_after_train=False``:
+    If you want to log objects after the training is finished use ``close_after_fit=False``:
 
     .. code-block:: python
 
@@ -135,7 +135,7 @@ class NeptuneLogger(LightningLoggerBase):
             "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.ai first.
-        offline_mode: Optional default False. If ``True`` no logs will be sent
+        offline_mode: Optional default ``False``. If ``True`` no logs will be sent
             to Neptune. Usually used for debug purposes.
         close_after_fit: Optional default ``True``. If ``False`` the experiment
             will not be closed after training and additional metrics,
@@ -167,6 +167,7 @@ class NeptuneLogger(LightningLoggerBase):
             Tags are displayed in the experiment’s Details section and can be viewed
             in the experiments view as a column.
     """
+
     def __init__(self,
                  api_key: Optional[str] = None,
                  project_name: Optional[str] = None,
@@ -188,24 +189,20 @@ def __init__(self,
         self.params = params
         self.properties = properties
         self.tags = tags
-        self._experiment = None
         self._kwargs = kwargs
+        self._experiment_id = None
+        self._experiment = self._create_or_get_experiment()
 
-        if offline_mode:
-            self.mode = 'offline'
-            neptune.init(project_qualified_name='dry-run/project',
-                         backend=neptune.OfflineBackend())
-        else:
-            self.mode = 'online'
-            neptune.init(api_token=self.api_key,
-                         project_qualified_name=self.project_name)
-
-        log.info(f'NeptuneLogger was initialized in {self.mode} mode')
+        log.info(f'NeptuneLogger will work in {"offline" if self.offline_mode else "online"} mode')
 
     def __getstate__(self):
         state = self.__dict__.copy()
-        # cannot be pickled
+
+        # Experiment cannot be pickled, and additionally its ID cannot be pickled in offline mode
         state['_experiment'] = None
+        if self.offline_mode:
+            state['_experiment_id'] = None
+
         return state
 
     @property
@@ -220,14 +217,11 @@ def experiment(self) -> Experiment:
 
         """
 
+        # Note that even though we initialize self._experiment in __init__,
+        # it may still end up being None after being pickled and un-pickled
         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)
+            self._experiment = self._create_or_get_experiment()
+
         return self._experiment
 
     @rank_zero_only
@@ -261,14 +255,14 @@ def finalize(self, status: str) -> None:
 
     @property
     def name(self) -> str:
-        if self.mode == 'offline':
+        if self.offline_mode:
             return 'offline-name'
         else:
             return self.experiment.name
 
     @property
     def version(self) -> str:
-        if self.mode == 'offline':
+        if self.offline_mode:
             return 'offline-id-1234'
         else:
             return self.experiment.id
@@ -363,3 +357,24 @@ def append_tags(self, tags: Union[str, Iterable[str]]) -> None:
         if str(tags) == tags:
             tags = [tags]  # make it as an iterable is if it is not yet
         self.experiment.append_tags(*tags)
+
+    def _create_or_get_experiment(self):
+        if self.offline_mode:
+            project = neptune.Session(backend=neptune.OfflineBackend()).get_project('dry-run/project')
+        else:
+            session = neptune.Session.with_default_backend(api_token=self.api_key)
+            project = session.get_project(self.project_name)
+
+        if self._experiment_id is None:
+            exp = project.create_experiment(
+                name=self.experiment_name,
+                params=self.params,
+                properties=self.properties,
+                tags=self.tags,
+                upload_source_files=self.upload_source_files,
+                **self._kwargs)
+        else:
+            exp = project.get_experiments(id=self._experiment_id)[0]
+
+        self._experiment_id = exp.id
+        return exp


Test patch:
diff --git a/tests/loggers/test_neptune.py b/tests/loggers/test_neptune.py
--- a/tests/loggers/test_neptune.py
+++ b/tests/loggers/test_neptune.py
@@ -10,53 +10,68 @@
 
 @patch('pytorch_lightning.loggers.neptune.neptune')
 def test_neptune_online(neptune):
-    logger = NeptuneLogger(api_key='test', offline_mode=False, project_name='project')
-    neptune.init.assert_called_once_with(api_token='test', project_qualified_name='project')
+    logger = NeptuneLogger(api_key='test', project_name='project')
 
-    assert logger.name == neptune.create_experiment().name
-    assert logger.version == neptune.create_experiment().id
+    created_experiment = neptune.Session.with_default_backend().get_project().create_experiment()
+
+    # It's important to check if the internal variable _experiment was initialized in __init__.
+    # Calling logger.experiment would cause a side-effect of initializing _experiment,
+    # if it wasn't already initialized.
+    assert logger._experiment == created_experiment
+    assert logger.name == created_experiment.name
+    assert logger.version == created_experiment.id
 
 
 @patch('pytorch_lightning.loggers.neptune.neptune')
-def test_neptune_additional_methods(neptune):
+def test_neptune_offline(neptune):
     logger = NeptuneLogger(offline_mode=True)
 
+    neptune.Session.assert_called_once_with(backend=neptune.OfflineBackend())
+    assert logger.experiment == neptune.Session().get_project().create_experiment()
+
+
+@patch('pytorch_lightning.loggers.neptune.neptune')
+def test_neptune_additional_methods(neptune):
+    logger = NeptuneLogger(api_key='test', project_name='project')
+
+    created_experiment = neptune.Session.with_default_backend().get_project().create_experiment()
+
     logger.log_metric('test', torch.ones(1))
-    neptune.create_experiment().log_metric.assert_called_once_with('test', torch.ones(1))
-    neptune.create_experiment().log_metric.reset_mock()
+    created_experiment.log_metric.assert_called_once_with('test', torch.ones(1))
+    created_experiment.log_metric.reset_mock()
 
     logger.log_metric('test', 1.0)
-    neptune.create_experiment().log_metric.assert_called_once_with('test', 1.0)
-    neptune.create_experiment().log_metric.reset_mock()
+    created_experiment.log_metric.assert_called_once_with('test', 1.0)
+    created_experiment.log_metric.reset_mock()
 
     logger.log_metric('test', 1.0, step=2)
-    neptune.create_experiment().log_metric.assert_called_once_with('test', x=2, y=1.0)
-    neptune.create_experiment().log_metric.reset_mock()
+    created_experiment.log_metric.assert_called_once_with('test', x=2, y=1.0)
+    created_experiment.log_metric.reset_mock()
 
     logger.log_text('test', 'text')
-    neptune.create_experiment().log_metric.assert_called_once_with('test', 'text')
-    neptune.create_experiment().log_metric.reset_mock()
+    created_experiment.log_metric.assert_called_once_with('test', 'text')
+    created_experiment.log_metric.reset_mock()
 
     logger.log_image('test', 'image file')
-    neptune.create_experiment().log_image.assert_called_once_with('test', 'image file')
-    neptune.create_experiment().log_image.reset_mock()
+    created_experiment.log_image.assert_called_once_with('test', 'image file')
+    created_experiment.log_image.reset_mock()
 
     logger.log_image('test', 'image file', step=2)
-    neptune.create_experiment().log_image.assert_called_once_with('test', x=2, y='image file')
-    neptune.create_experiment().log_image.reset_mock()
+    created_experiment.log_image.assert_called_once_with('test', x=2, y='image file')
+    created_experiment.log_image.reset_mock()
 
     logger.log_artifact('file')
-    neptune.create_experiment().log_artifact.assert_called_once_with('file', None)
+    created_experiment.log_artifact.assert_called_once_with('file', None)
 
     logger.set_property('property', 10)
-    neptune.create_experiment().set_property.assert_called_once_with('property', 10)
+    created_experiment.set_property.assert_called_once_with('property', 10)
 
     logger.append_tags('one tag')
-    neptune.create_experiment().append_tags.assert_called_once_with('one tag')
-    neptune.create_experiment().append_tags.reset_mock()
+    created_experiment.append_tags.assert_called_once_with('one tag')
+    created_experiment.append_tags.reset_mock()
 
     logger.append_tags(['two', 'tags'])
-    neptune.create_experiment().append_tags.assert_called_once_with('two', 'tags')
+    created_experiment.append_tags.assert_called_once_with('two', 'tags')
 
 
 def test_neptune_leave_open_experiment_after_fit(tmpdir):