From 5a9afb11cc270cae5e59048dd78a894e90170ab3 Mon Sep 17 00:00:00 2001
From: Ir1dXD <sirius.caffrey@gmail.com>
Date: Tue, 5 Nov 2019 21:43:21 +0800
Subject: [PATCH 1/6] change print to logging (#457)

* change print to logging

* always use logging.info

* use f-strings

* update code style

* set logging configs

* remove unused code
---
 README.md                                     |  4 +-
 docs/examples/Examples.md                     |  2 +-
 .../lightning_module_template.py              |  7 ++--
 pytorch_lightning/callbacks/pt_callbacks.py   | 40 ++++++++++---------
 pytorch_lightning/root_module/memory.py       |  3 +-
 pytorch_lightning/root_module/root_module.py  |  3 +-
 pytorch_lightning/trainer/amp_mixin.py        |  3 +-
 pytorch_lightning/trainer/ddp_mixin.py        |  5 ++-
 pytorch_lightning/trainer/trainer.py          |  6 ++-
 pytorch_lightning/trainer/trainer_io.py       | 17 ++++----
 .../trainer/training_tricks_mixin.py          |  4 +-
 tests/test_a_restore_models.py                |  3 +-
 12 files changed, 56 insertions(+), 41 deletions(-)

diff --git a/README.md b/README.md
index b91857a7..f0cd30db 100644
--- a/README.md
+++ b/README.md
@@ -164,8 +164,8 @@ trainer = Trainer(max_nb_epochs=1, train_percent_check=0.1)
 trainer.fit(model)
 
 # view tensorboard logs 
-print('View tensorboard logs by running\ntensorboard --logdir %s' % os.getcwd())
-print('and going to http://localhost:6006 on your browser')
+logging.info(f'View tensorboard logs by running\ntensorboard --logdir {os.getcwd()}')
+logging.info('and going to http://localhost:6006 on your browser')
 ```    
 
 When you're all done you can even run the test set separately.   
diff --git a/docs/examples/Examples.md b/docs/examples/Examples.md
index 5deda159..49e83e70 100644
--- a/docs/examples/Examples.md
+++ b/docs/examples/Examples.md
@@ -119,7 +119,7 @@ def optimize_on_cluster(hyperparams):
     job_display_name = job_display_name[0:3]
 
     # run hopt
-    print('submitting jobs...')
+    logging.info('submitting jobs...')
     cluster.optimize_parallel_cluster_gpu(
         main,
         nb_trials=hyperparams.nb_hopt_trials,
diff --git a/pl_examples/basic_examples/lightning_module_template.py b/pl_examples/basic_examples/lightning_module_template.py
index 6b58ab46..8a892d10 100644
--- a/pl_examples/basic_examples/lightning_module_template.py
+++ b/pl_examples/basic_examples/lightning_module_template.py
@@ -2,6 +2,7 @@
 Example template for defining a system
 """
 import os
+import logging
 from argparse import ArgumentParser
 from collections import OrderedDict
 
@@ -214,17 +215,17 @@ class LightningTemplateModel(LightningModule):
 
     @pl.data_loader
     def train_dataloader(self):
-        print('training data loader called')
+        logging.info('training data loader called')
         return self.__dataloader(train=True)
 
     @pl.data_loader
     def val_dataloader(self):
-        print('val data loader called')
+        logging.info('val data loader called')
         return self.__dataloader(train=False)
 
     @pl.data_loader
     def test_dataloader(self):
-        print('test data loader called')
+        logging.info('test data loader called')
         return self.__dataloader(train=False)
 
     @staticmethod
diff --git a/pytorch_lightning/callbacks/pt_callbacks.py b/pytorch_lightning/callbacks/pt_callbacks.py
index d5d80d24..2c1553d4 100644
--- a/pytorch_lightning/callbacks/pt_callbacks.py
+++ b/pytorch_lightning/callbacks/pt_callbacks.py
@@ -1,6 +1,7 @@
 import os
 import shutil
-
+import logging
+import warnings
 import numpy as np
 
 from pytorch_lightning.pt_overrides.override_data_parallel import LightningDistributedDataParallel
@@ -91,7 +92,7 @@ class EarlyStopping(Callback):
         self.stopped_epoch = 0
 
         if mode not in ['auto', 'min', 'max']:
-            print('EarlyStopping mode %s is unknown, fallback to auto mode.' % mode)
+            logging.info(f'EarlyStopping mode {mode} is unknown, fallback to auto mode.')
             mode = 'auto'
 
         if mode == 'min':
@@ -121,9 +122,10 @@ class EarlyStopping(Callback):
         current = logs.get(self.monitor)
         stop_training = False
         if current is None:
-            print('Early stopping conditioned on metric `%s` '
-                  'which is not available. Available metrics are: %s' %
-                  (self.monitor, ','.join(list(logs.keys()))), RuntimeWarning)
+            warnings.warn(
+                f'Early stopping conditioned on metric `{self.monitor}`'
+                f' which is not available. Available metrics are: {",".join(list(logs.keys()))}',
+                RuntimeWarning)
             stop_training = True
             return stop_training
 
@@ -141,7 +143,7 @@ class EarlyStopping(Callback):
 
     def on_train_end(self, logs=None):
         if self.stopped_epoch > 0 and self.verbose > 0:
-            print('Epoch %05d: early stopping' % (self.stopped_epoch + 1))
+            logging.info(f'Epoch {self.stopped_epoch + 1:05d}: early stopping')
 
 
 class ModelCheckpoint(Callback):
@@ -187,8 +189,9 @@ class ModelCheckpoint(Callback):
         self.prefix = prefix
 
         if mode not in ['auto', 'min', 'max']:
-            print('ModelCheckpoint mode %s is unknown, '
-                  'fallback to auto mode.' % (mode), RuntimeWarning)
+            warnings.warn(
+                f'ModelCheckpoint mode {mode} is unknown, '
+                'fallback to auto mode.', RuntimeWarning)
             mode = 'auto'
 
         if mode == 'min':
@@ -232,25 +235,26 @@ class ModelCheckpoint(Callback):
             if self.save_best_only:
                 current = logs.get(self.monitor)
                 if current is None:
-                    print('Can save best model only with %s available,'
-                          ' skipping.' % (self.monitor), RuntimeWarning)
+                    warnings.warn(
+                        f'Can save best model only with {self.monitor} available,'
+                        ' skipping.', RuntimeWarning)
                 else:
                     if self.monitor_op(current, self.best):
                         if self.verbose > 0:
-                            print('\nEpoch %05d: %s improved from %0.5f to %0.5f,'
-                                  ' saving model to %s'
-                                  % (epoch + 1, self.monitor, self.best,
-                                     current, filepath))
+                            logging.info(
+                                f'\nEpoch {epoch + 1:05d}: {self.monitor} improved'
+                                f' from {self.best:0.5f} to {current:0.5f},',
+                                f' saving model to {filepath}')
                         self.best = current
                         self.save_model(filepath, overwrite=True)
 
                     else:
                         if self.verbose > 0:
-                            print('\nEpoch %05d: %s did not improve' %
-                                  (epoch + 1, self.monitor))
+                            logging.info(
+                                f'\nEpoch {epoch + 1:05d}: {self.monitor} did not improve')
             else:
                 if self.verbose > 0:
-                    print('\nEpoch %05d: saving model to %s' % (epoch + 1, filepath))
+                    logging.info(f'\nEpoch {epoch + 1:05d}: saving model to {filepath}')
                 self.save_model(filepath, overwrite=False)
 
 
@@ -291,6 +295,6 @@ if __name__ == '__main__':
     losses = [10, 9, 8, 8, 6, 4.3, 5, 4.4, 2.8, 2.5]
     for i, loss in enumerate(losses):
         should_stop = c.on_epoch_end(i, logs={'val_loss': loss})
-        print(loss)
+        logging.info(loss)
         if should_stop:
             break
diff --git a/pytorch_lightning/root_module/memory.py b/pytorch_lightning/root_module/memory.py
index d16d2659..1201aab9 100644
--- a/pytorch_lightning/root_module/memory.py
+++ b/pytorch_lightning/root_module/memory.py
@@ -8,6 +8,7 @@ import subprocess
 import numpy as np
 import pandas as pd
 import torch
+import logging
 
 
 class ModelSummary(object):
@@ -166,7 +167,7 @@ def print_mem_stack():  # pragma: no cover
     for obj in gc.get_objects():
         try:
             if torch.is_tensor(obj) or (hasattr(obj, 'data') and torch.is_tensor(obj.data)):
-                print(type(obj), obj.size())
+                logging.info(type(obj), obj.size())
         except Exception:
             pass
 
diff --git a/pytorch_lightning/root_module/root_module.py b/pytorch_lightning/root_module/root_module.py
index ab5275dc..b25ab9e4 100644
--- a/pytorch_lightning/root_module/root_module.py
+++ b/pytorch_lightning/root_module/root_module.py
@@ -10,6 +10,7 @@ from pytorch_lightning.root_module.hooks import ModelHooks
 from pytorch_lightning.root_module.memory import ModelSummary
 from pytorch_lightning.root_module.model_saving import ModelIO
 from pytorch_lightning.trainer.trainer_io import load_hparams_from_tags_csv
+import logging
 
 
 class LightningModule(GradInformation, ModelIO, ModelHooks):
@@ -240,7 +241,7 @@ class LightningModule(GradInformation, ModelIO, ModelHooks):
 
     def summarize(self, mode):
         model_summary = ModelSummary(self, mode=mode)
-        print(model_summary)
+        logging.info(model_summary)
 
     def freeze(self):
         for param in self.parameters():
diff --git a/pytorch_lightning/trainer/amp_mixin.py b/pytorch_lightning/trainer/amp_mixin.py
index d0790242..42251bb3 100644
--- a/pytorch_lightning/trainer/amp_mixin.py
+++ b/pytorch_lightning/trainer/amp_mixin.py
@@ -4,6 +4,7 @@ try:
     APEX_AVAILABLE = True
 except ImportError:
     APEX_AVAILABLE = False
+import logging
 
 
 class TrainerAMPMixin(object):
@@ -11,7 +12,7 @@ class TrainerAMPMixin(object):
     def init_amp(self, use_amp):
         self.use_amp = use_amp and APEX_AVAILABLE
         if self.use_amp:
-            print('using 16bit precision')
+            logging.info('using 16bit precision')
 
         if use_amp and not APEX_AVAILABLE:  # pragma: no cover
             msg = """
diff --git a/pytorch_lightning/trainer/ddp_mixin.py b/pytorch_lightning/trainer/ddp_mixin.py
index 36345b74..9f6c1ad4 100644
--- a/pytorch_lightning/trainer/ddp_mixin.py
+++ b/pytorch_lightning/trainer/ddp_mixin.py
@@ -1,6 +1,7 @@
 import os
 import re
 import warnings
+import logging
 
 import torch
 import torch.distributed as dist
@@ -59,7 +60,7 @@ class TrainerDDPMixin(object):
                 'To silence this warning set distributed_backend=ddp'
             warnings.warn(w)
 
-        print('gpu available: {}, used: {}'.format(torch.cuda.is_available(), self.on_gpu))
+        logging.info(f'gpu available: {torch.cuda.is_available()}, used: {self.on_gpu}')
 
     def configure_slurm_ddp(self, nb_gpu_nodes):
         self.is_slurm_managing_tasks = False
@@ -107,7 +108,7 @@ class TrainerDDPMixin(object):
                 gpu_str = ','.join([str(x) for x in data_parallel_device_ids])
                 os.environ["CUDA_VISIBLE_DEVICES"] = gpu_str
 
-        print(f'VISIBLE GPUS: {os.environ["CUDA_VISIBLE_DEVICES"]}')
+        logging.info(f'VISIBLE GPUS: {os.environ["CUDA_VISIBLE_DEVICES"]}')
 
     def ddp_train(self, gpu_nb, model):
         """
diff --git a/pytorch_lightning/trainer/trainer.py b/pytorch_lightning/trainer/trainer.py
index c53aeb7b..eebba265 100644
--- a/pytorch_lightning/trainer/trainer.py
+++ b/pytorch_lightning/trainer/trainer.py
@@ -4,6 +4,7 @@ The trainer handles all the logic for running a val loop, training loop, distrib
 
 import os
 import warnings
+import logging
 
 import torch
 import torch.distributed as dist
@@ -148,7 +149,7 @@ class Trainer(TrainerIOMixin,
             Running in fast_dev_run mode: will run a full train,
             val loop using a single batch
             '''
-            print(m)
+            logging.info(m)
 
         # set default save path if user didn't provide one
         self.default_save_path = default_save_path
@@ -234,6 +235,9 @@ class Trainer(TrainerIOMixin,
         self.amp_level = amp_level
         self.init_amp(use_amp)
 
+        # set logging options
+        logging.basicConfig(level=logging.INFO)
+
     @property
     def slurm_job_id(self):
         try:
diff --git a/pytorch_lightning/trainer/trainer_io.py b/pytorch_lightning/trainer/trainer_io.py
index dd38c051..11bbb1d5 100644
--- a/pytorch_lightning/trainer/trainer_io.py
+++ b/pytorch_lightning/trainer/trainer_io.py
@@ -3,6 +3,7 @@ import re
 import signal
 import warnings
 from subprocess import call
+import logging
 
 import torch
 import torch.distributed as dist
@@ -87,7 +88,7 @@ class TrainerIOMixin(object):
         if last_ckpt_name is not None:
             last_ckpt_path = os.path.join(self.checkpoint_callback.filepath, last_ckpt_name)
             self.restore(last_ckpt_path, self.on_gpu)
-            print(f'model and trainer restored from checkpoint: {last_ckpt_path}')
+            logging.info(f'model and trainer restored from checkpoint: {last_ckpt_path}')
             did_restore = True
 
         return did_restore
@@ -106,14 +107,14 @@ class TrainerIOMixin(object):
             pass
 
         if on_slurm:
-            print('set slurm handle signals')
+            logging.info('set slurm handle signals')
             signal.signal(signal.SIGUSR1, self.sig_handler)
             signal.signal(signal.SIGTERM, self.term_handler)
 
     def sig_handler(self, signum, frame):
         if self.proc_rank == 0:
             # save weights
-            print('handling SIGUSR1')
+            logging.info('handling SIGUSR1')
             self.hpc_save(self.weights_save_path, self.logger)
 
             # find job id
@@ -121,21 +122,21 @@ class TrainerIOMixin(object):
             cmd = 'scontrol requeue {}'.format(job_id)
 
             # requeue job
-            print('\nrequeing job {}...'.format(job_id))
+            logging.info('\nrequeing job {job_id}...')
             result = call(cmd, shell=True)
 
             # print result text
             if result == 0:
-                print('requeued exp ', job_id)
+                logging.info('requeued exp {job_id}')
             else:
-                print('requeue failed...')
+                logging.info('requeue failed...')
 
             # close experiment to avoid issues
             self.logger.close()
 
     def term_handler(self, signum, frame):
         # save
-        print("bypassing sigterm")
+        logging.info("bypassing sigterm")
 
     # --------------------
     # MODEL SAVE CHECKPOINT
@@ -328,7 +329,7 @@ class TrainerIOMixin(object):
         # call model hook
         model.on_hpc_load(checkpoint)
 
-        print(f'restored hpc model from: {filepath}')
+        logging.info(f'restored hpc model from: {filepath}')
 
     def max_ckpt_in_folder(self, path, name_key='ckpt_'):
         files = os.listdir(path)
diff --git a/pytorch_lightning/trainer/training_tricks_mixin.py b/pytorch_lightning/trainer/training_tricks_mixin.py
index 1efc9a17..836b6689 100644
--- a/pytorch_lightning/trainer/training_tricks_mixin.py
+++ b/pytorch_lightning/trainer/training_tricks_mixin.py
@@ -1,5 +1,5 @@
 import torch
-
+import logging
 from pytorch_lightning.callbacks import GradientAccumulationScheduler
 
 
@@ -14,7 +14,7 @@ class TrainerTrainingTricksMixin(object):
         model = self.get_model()
         for param in model.parameters():
             if torch.isnan(param.grad.float()).any():
-                print(param, param.grad)
+                logging.info(param, param.grad)
 
     def configure_accumulated_gradients(self, accumulate_grad_batches):
         self.accumulate_grad_batches = None
diff --git a/tests/test_a_restore_models.py b/tests/test_a_restore_models.py
index 5b747aec..22ef9f11 100644
--- a/tests/test_a_restore_models.py
+++ b/tests/test_a_restore_models.py
@@ -1,4 +1,5 @@
 import os
+import logging
 
 import pytest
 import torch
@@ -44,7 +45,7 @@ def test_running_test_pretrained_model_ddp():
     result = trainer.fit(model)
 
     exp = logger.experiment
-    print(os.listdir(exp.get_data_path(exp.name, exp.version)))
+    logging.info(os.listdir(exp.get_data_path(exp.name, exp.version)))
 
     # correct result and ok accuracy
     assert result == 1, 'training failed to complete'

From 32dd803b1ebed02d3ca05705b719880b57b8830a Mon Sep 17 00:00:00 2001
From: Yongrae Jo <dreamgonfly@gmail.com>
Date: Tue, 5 Nov 2019 22:55:44 +0900
Subject: [PATCH 2/6] Fix min_max gpu memory logging bug (#453)

* #452 Fix ValueError

* #452 Use subprocess.run

* #452 Simplify code for gpu_memory_map

* #452 Simplify code for min max memory

* #452 Add test for get_memory_profile

* #452 Use os.sep

* #452 Use os.linesep
---
 pytorch_lightning/root_module/memory.py | 35 ++++++++++---------------
 tests/test_gpu_models.py                |  2 +-
 2 files changed, 15 insertions(+), 22 deletions(-)

diff --git a/pytorch_lightning/root_module/memory.py b/pytorch_lightning/root_module/memory.py
index 1201aab9..362688eb 100644
--- a/pytorch_lightning/root_module/memory.py
+++ b/pytorch_lightning/root_module/memory.py
@@ -3,6 +3,7 @@ Generates a summary of a model's layers and dimensionality
 '''
 
 import gc
+import os
 import subprocess
 
 import numpy as np
@@ -199,19 +200,10 @@ def get_memory_profile(mode):
     memory_map = get_gpu_memory_map()
 
     if mode == 'min_max':
-        min_mem = 1000000
-        min_k = None
-        max_mem = 0
-        max_k = None
-        for k, v in memory_map:
-            if v > max_mem:
-                max_mem = v
-                max_k = k
-            if v < min_mem:
-                min_mem = v
-                min_k = k
+        min_index, min_memory = min(memory_map.items(), key=lambda item: item[1])
+        max_index, max_memory = max(memory_map.items(), key=lambda item: item[1])
 
-        memory_map = {min_k: min_mem, max_k: max_mem}
+        memory_map = {min_index: min_memory, max_index: max_memory}
 
     return memory_map
 
@@ -225,17 +217,18 @@ def get_gpu_memory_map():
         Keys are device ids as integers.
         Values are memory usage as integers in MB.
     """
-    result = subprocess.check_output(
+    result = subprocess.run(
         [
-            'nvidia-smi', '--query-gpu=memory.used',
-            '--format=csv,nounits,noheader'
-        ], encoding='utf-8')
+            'nvidia-smi',
+            '--query-gpu=memory.used',
+            '--format=csv,nounits,noheader',
+        ],
+        encoding='utf-8',
+        capture_output=True,
+        check=True)
     # Convert lines into a dictionary
-    gpu_memory = [int(x) for x in result.strip().split('\n')]
-    gpu_memory_map = {}
-    for k, v in zip(range(len(gpu_memory)), gpu_memory):
-        k = f'gpu_{k}'
-        gpu_memory_map[k] = v
+    gpu_memory = [int(x) for x in result.stdout.strip().split(os.linesep)]
+    gpu_memory_map = {f'gpu_{index}': memory for index, memory in enumerate(gpu_memory)}
     return gpu_memory_map
 
 
diff --git a/tests/test_gpu_models.py b/tests/test_gpu_models.py
index 118a3f2a..17f8a8fd 100644
--- a/tests/test_gpu_models.py
+++ b/tests/test_gpu_models.py
@@ -224,7 +224,7 @@ def test_multi_gpu_model_dp():
     testing_utils.run_gpu_model_test(trainer_options, model, hparams)
 
     # test memory helper functions
-    memory.get_gpu_memory_map()
+    memory.get_memory_profile('min_max')
 
 
 def test_ddp_sampler_error():

From 8fbaccddaed6dc574e02fbe5d2f3dae161336785 Mon Sep 17 00:00:00 2001
From: William Falcon <waf2107@columbia.edu>
Date: Tue, 5 Nov 2019 09:14:33 -0500
Subject: [PATCH 3/6] added eval and train for redundancy (#464)

---
 pytorch_lightning/root_module/root_module.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/pytorch_lightning/root_module/root_module.py b/pytorch_lightning/root_module/root_module.py
index b25ab9e4..fc8d62c5 100644
--- a/pytorch_lightning/root_module/root_module.py
+++ b/pytorch_lightning/root_module/root_module.py
@@ -247,6 +247,10 @@ class LightningModule(GradInformation, ModelIO, ModelHooks):
         for param in self.parameters():
             param.requires_grad = False
 
+        self.eval()
+
     def unfreeze(self):
         for param in self.parameters():
             param.requires_grad = True
+
+        self.train()

From 3e38005a61f5472d9f2938e745e7ac958ff4916a Mon Sep 17 00:00:00 2001
From: William Falcon <waf2107@columbia.edu>
Date: Tue, 5 Nov 2019 10:01:52 -0500
Subject: [PATCH 4/6] Ddp2 fix (#448)

* added training_end

* added training_end

* added training_end

* added training_end

* added training_end

* added training_end

* added training_end

* added training_end

* added training_end

* added training_end

* added training_end

* added training_end

* allow ddp and apex to be configured

* allow ddp and apex to be configured

* bananas

* bananas

* bananas

* bananas

* bananas

* bananas

* bananas

* bananas

* bananas

* bananas

* bananas

* bananas

* bananas

* bananas

* bananas

* added eval and train for redundancy

* added eval and train for redundancy

* added training_end

* added training_end

* added training_end

* added training_end

* added training_end

* added training_end

* added training_end

* added training_end

* added training_end

* added training_end

* added training_end

* added training_end

* allow ddp and apex to be configured

* allow ddp and apex to be configured

* bananas

* bananas

* bananas

* bananas

* bananas

* bananas

* bananas

* bananas

* bananas

* bananas

* bananas

* bananas

* bananas

* bananas

* bananas

* added eval and train for redundancy

* added eval and train for redundancy
---
 README.md                                     |  1 +
 .../RequiredTrainerInterface.md               | 84 +++++++++++++++++
 docs/Trainer/hooks.md                         | 89 +++++++++++++++++++
 docs/Trainer/index.md                         |  1 +
 docs/index.md                                 |  1 +
 pytorch_lightning/root_module/root_module.py  | 80 ++++++++++++++++-
 pytorch_lightning/trainer/ddp_mixin.py        | 54 ++---------
 pytorch_lightning/trainer/dp_mixin.py         |  4 +-
 pytorch_lightning/trainer/logging_mixin.py    |  4 +
 pytorch_lightning/trainer/train_loop_mixin.py | 29 +++---
 10 files changed, 284 insertions(+), 63 deletions(-)

diff --git a/README.md b/README.md
index f0cd30db..0db819f9 100644
--- a/README.md
+++ b/README.md
@@ -294,6 +294,7 @@ Lightning also adds a text column with all the hyperparameters for this experime
 
 #### Distributed training    
 
+- [Implement Your Own Distributed (DDP) training](https://williamfalcon.github.io/pytorch-lightning/Trainer/hooks/#init_ddp_connection)
 - [16-bit mixed precision](https://williamfalcon.github.io/pytorch-lightning/Trainer/Distributed%20training/#16-bit-mixed-precision)
 - [Multi-GPU](https://williamfalcon.github.io/pytorch-lightning/Trainer/Distributed%20training/#Multi-GPU)
 - [Multi-node](https://williamfalcon.github.io/pytorch-lightning/Trainer/Distributed%20training/#Multi-node)
diff --git a/docs/LightningModule/RequiredTrainerInterface.md b/docs/LightningModule/RequiredTrainerInterface.md
index 53bfc821..7881dde4 100644
--- a/docs/LightningModule/RequiredTrainerInterface.md
+++ b/docs/LightningModule/RequiredTrainerInterface.md
@@ -15,6 +15,7 @@ Otherwise, to Define a Lightning Module, implement the following methods:
 
 **Optional**:   
 
+- [training_end](RequiredTrainerInterface.md#training_end)      
 - [validation_step](RequiredTrainerInterface.md#validation_step)    
 - [validation_end](RequiredTrainerInterface.md#validation_end) 
 - [test_step](RequiredTrainerInterface.md#test_step)    
@@ -178,6 +179,89 @@ def training_step(self, batch, batch_nb, hiddens):
 You can also return a -1 instead of a dict to stop the current loop. This is useful if you want to
 break out of the current training epoch early.
 
+---   
+### training_end
+
+``` {.python}
+def training_end(self, train_step_outputs)
+```
+In certain cases (dp, ddp2), you might want to use all outputs of every process to do something.
+For instance, if using negative samples, you could run a batch via dp and use ALL the outputs 
+for a single softmax across the full batch (ie: the denominator would use the full batch).
+
+In this case you should define training_end to perform those calculations.
+
+
+**Params**    
+
+| Param  | description  |
+|---|---|
+|  outputs | What you return in training_step. 
+
+**Return**   
+
+Dictionary or OrderedDict   
+
+| key  | value  | is required |
+|---|---|---|
+|  loss | tensor scalar  | Y |
+|  progress_bar | Dict for progress bar display. Must have only tensors  | N |
+|  log | Dict of metrics to add to logger. Must have only tensors (no images, etc)  | N |
+
+
+**Example**
+
+``` {.python}
+# WITHOUT training_end
+# if used in DP or DDP2, this batch is 1/nb_gpus large
+def training_step(self, batch, batch_nb):
+    # batch is 1/nb_gpus big
+    x, y = batch
+    
+    out = self.forward(x)
+    loss = self.softmax(out)
+    loss = nce_loss(loss)
+    return {'loss': loss}
+
+# --------------
+# with training_end to do softmax over the full batch
+def training_step(self, batch, batch_nb):
+    # batch is 1/nb_gpus big
+    x, y = batch
+    
+    out = self.forward(x)
+    return {'out': out}
+
+def training_end(self, outputs):
+    # this out is now the full size of the batch
+    out = outputs['out']
+
+    # this softmax now uses the full batch size
+    loss = self.softmax(out)
+    loss = nce_loss(loss)
+    return {'loss': loss}
+```    
+
+If you define multiple optimizers, this step will also be called with an additional ```optimizer_idx``` param.    
+``` {.python}
+# Multiple optimizers (ie: GANs)     
+def training_step(self, batch, batch_nb, optimizer_idx):
+    if optimizer_idx == 0:
+        # do training_step with encoder
+    if optimizer_idx == 1:
+        # do training_step with decoder    
+```    
+
+If you add truncated back propagation through time you will also get an additional argument with the hidden states of the previous step.     
+``` {.python}
+# Truncated back-propagation through time   
+def training_step(self, batch, batch_nb, hiddens):
+    # hiddens are the hiddens from the previous truncated backprop step
+```    
+
+You can also return a -1 instead of a dict to stop the current loop. This is useful if you want to
+break out of the current training epoch early.
+
 --- 
 ### train_dataloader 
 
diff --git a/docs/Trainer/hooks.md b/docs/Trainer/hooks.md
index 96c41e2b..5f1baf84 100644
--- a/docs/Trainer/hooks.md
+++ b/docs/Trainer/hooks.md
@@ -175,3 +175,92 @@ def tbptt_split_batch(self, batch, split_size):
 
   return splits
 ```
+
+---
+#### configure_apex
+Overwrite to define your own Apex implementation init.
+
+```python
+def configure_apex(self, amp, model, optimizers, amp_level):
+    """
+    Override to init AMP your own way
+    Must return a model and list of optimizers
+    :param amp:
+    :param model:
+    :param optimizers:
+    :param amp_level:
+    :return: Apex wrapped model and optimizers
+    """
+    model, optimizers = amp.initialize(
+        model, optimizers, opt_level=amp_level,
+    )
+
+    return model, optimizers
+```
+
+---
+#### configure_ddp 
+Overwrite to define your own DDP implementation init.
+The only requirement is that:
+1. On a validation batch the call goes to model.validation_step.   
+2. On a training batch the call goes to model.training_step.   
+3. On a testing batch, the call goes to model.test_step
+
+```python
+def configure_ddp(self, model, device_ids):
+    """
+    Override to init DDP in a different way or use your own wrapper.
+    Must return model.
+    :param model:
+    :param device_ids:
+    :return: DDP wrapped model
+    """
+    # Lightning DDP simply routes to test_step, val_step, etc...
+    model = LightningDistributedDataParallel(
+        model,
+        device_ids=device_ids,
+        find_unused_parameters=True
+    )
+    return model
+```
+
+---   
+#### init_ddp_connection   
+Override to init DDP in your own way.   
+
+```python
+def init_ddp_connection(self):
+    """
+    Connect all procs in the world using the env:// init
+    Use the first node as the root address
+    """
+
+    # use slurm job id for the port number
+    # guarantees unique ports across jobs from same grid search
+    try:
+        # use the last 4 numbers in the job id as the id
+        default_port = os.environ['SLURM_JOB_ID']
+        default_port = default_port[-4:]
+
+        # all ports should be in the 10k+ range
+        default_port = int(default_port) + 15000
+
+    except Exception as e:
+        default_port = 12910
+
+    # if user gave a port number, use that one instead
+    try:
+        default_port = os.environ['MASTER_PORT']
+    except Exception:
+        os.environ['MASTER_PORT'] = str(default_port)
+
+    # figure out the root node addr
+    try:
+        root_node = os.environ['SLURM_NODELIST'].split(' ')[0]
+    except Exception:
+        root_node = '127.0.0.2'
+
+    root_node = self.trainer.resolve_root_node_address(root_node)
+    os.environ['MASTER_ADDR'] = root_node
+    dist.init_process_group('nccl', rank=self.proc_rank, world_size=self.world_size)
+```
diff --git a/docs/Trainer/index.md b/docs/Trainer/index.md
index 13df0539..1c60786c 100644
--- a/docs/Trainer/index.md
+++ b/docs/Trainer/index.md
@@ -42,6 +42,7 @@ But of course the fun is in all the advanced things it can do:
 
 **Distributed training**    
 
+- [Implement Your Own Distributed (DDP) training](https://williamfalcon.github.io/pytorch-lightning/Trainer/hooks/#init_ddp_connection)
 - [16-bit mixed precision](https://williamfalcon.github.io/pytorch-lightning/Trainer/Distributed%20training/#16-bit-mixed-precision)
 - [Multi-GPU](https://williamfalcon.github.io/pytorch-lightning/Trainer/Distributed%20training/#Multi-GPU)
 - [Multi-node](https://williamfalcon.github.io/pytorch-lightning/Trainer/Distributed%20training/#Multi-node)
diff --git a/docs/index.md b/docs/index.md
index 06a76fec..b48b0d53 100644
--- a/docs/index.md
+++ b/docs/index.md
@@ -99,6 +99,7 @@ Notice a few things about this flow:
 
 ###### Distributed training    
 
+- [Implement Your Own Distributed (DDP) training](https://williamfalcon.github.io/pytorch-lightning/Trainer/hooks/#init_ddp_connection)
 - [16-bit mixed precision](https://williamfalcon.github.io/pytorch-lightning/Trainer/Distributed%20training/#16-bit-mixed-precision)
 - [Multi-GPU](https://williamfalcon.github.io/pytorch-lightning/Trainer/Distributed%20training/#Multi-GPU)
 - [Multi-node](https://williamfalcon.github.io/pytorch-lightning/Trainer/Distributed%20training/#Multi-node)
diff --git a/pytorch_lightning/root_module/root_module.py b/pytorch_lightning/root_module/root_module.py
index fc8d62c5..c5703e8b 100644
--- a/pytorch_lightning/root_module/root_module.py
+++ b/pytorch_lightning/root_module/root_module.py
@@ -1,8 +1,10 @@
+import os
 import warnings
 import collections
 from argparse import Namespace
 
 import torch
+import torch.distributed as dist
 
 from pytorch_lightning.root_module.decorators import data_loader
 from pytorch_lightning.root_module.grads import GradInformation
@@ -11,6 +13,7 @@ from pytorch_lightning.root_module.memory import ModelSummary
 from pytorch_lightning.root_module.model_saving import ModelIO
 from pytorch_lightning.trainer.trainer_io import load_hparams_from_tags_csv
 import logging
+from pytorch_lightning.pt_overrides.override_data_parallel import LightningDistributedDataParallel
 
 
 class LightningModule(GradInformation, ModelIO, ModelHooks):
@@ -48,10 +51,19 @@ class LightningModule(GradInformation, ModelIO, ModelHooks):
         return loss, dict with metrics for tqdm
         :param called with batch, batch_nb
         additional: optimizer_i if multiple optimizers used
-        :return:
+        :return: dict with loss key and optional log, progress keys
+                if implementing training_step, return whatever you need in that step
         """
         raise NotImplementedError
 
+    def training_end(self, *args, **kwargs):
+        """
+        return loss, dict with metrics for tqdm
+        :param called with outputs of training_step
+        :return: dict with loss key and optional log, progress keys
+        """
+        pass
+
     def validation_step(self, *args, **kwargs):
         """
         return whatever outputs will need to be aggregated in validation_end
@@ -90,6 +102,72 @@ class LightningModule(GradInformation, ModelIO, ModelHooks):
         """
         pass
 
+    def configure_ddp(self, model, device_ids):
+        """
+        Override to init DDP in a different way or use your own wrapper.
+        Must return model.
+        :param model:
+        :param device_ids:
+        :return: DDP wrapped model
+        """
+        model = LightningDistributedDataParallel(
+            model,
+            device_ids=device_ids,
+            find_unused_parameters=True
+        )
+        return model
+
+    def init_ddp_connection(self, proc_rank, world_size):
+        """
+        Connect all procs in the world using the env:// init
+        Use the first node as the root address
+        """
+
+        # use slurm job id for the port number
+        # guarantees unique ports across jobs from same grid search
+        try:
+            # use the last 4 numbers in the job id as the id
+            default_port = os.environ['SLURM_JOB_ID']
+            default_port = default_port[-4:]
+
+            # all ports should be in the 10k+ range
+            default_port = int(default_port) + 15000
+
+        except Exception as e:
+            default_port = 12910
+
+        # if user gave a port number, use that one instead
+        try:
+            default_port = os.environ['MASTER_PORT']
+        except Exception:
+            os.environ['MASTER_PORT'] = str(default_port)
+
+        # figure out the root node addr
+        try:
+            root_node = os.environ['SLURM_NODELIST'].split(' ')[0]
+        except Exception:
+            root_node = '127.0.0.2'
+
+        root_node = self.trainer.resolve_root_node_address(root_node)
+        os.environ['MASTER_ADDR'] = root_node
+        dist.init_process_group('nccl', rank=proc_rank, world_size=world_size)
+
+    def configure_apex(self, amp, model, optimizers, amp_level):
+        """
+        Override to init AMP your own way
+        Must return a model and list of optimizers
+        :param amp:
+        :param model:
+        :param optimizers:
+        :param amp_level:
+        :return: Apex wrapped model and optimizers
+        """
+        model, optimizers = amp.initialize(
+            model, optimizers, opt_level=amp_level,
+        )
+
+        return model, optimizers
+
     def configure_optimizers(self):
         """
         Return a list of optimizers and a list of schedulers (could be empty)
diff --git a/pytorch_lightning/trainer/ddp_mixin.py b/pytorch_lightning/trainer/ddp_mixin.py
index 9f6c1ad4..0653fa58 100644
--- a/pytorch_lightning/trainer/ddp_mixin.py
+++ b/pytorch_lightning/trainer/ddp_mixin.py
@@ -4,9 +4,7 @@ import warnings
 import logging
 
 import torch
-import torch.distributed as dist
 
-from pytorch_lightning.pt_overrides.override_data_parallel import LightningDistributedDataParallel
 from pytorch_lightning.utilities.debugging import MisconfigurationException
 
 try:
@@ -145,7 +143,8 @@ class TrainerDDPMixin(object):
         # set up server using proc 0's ip address
         # try to init for 20 times at max in case ports are taken
         # where to store ip_table
-        self.__init_tcp_connection()
+        model.trainer = self
+        model.init_ddp_connection(self.proc_rank, self.world_size)
 
         # CHOOSE OPTIMIZER
         # allow for lr schedulers as well
@@ -167,9 +166,7 @@ class TrainerDDPMixin(object):
         # run through amp wrapper before going to distributed DP
         if self.use_amp:
             # An example
-            model, optimizers = amp.initialize(
-                model, self.optimizers, opt_level=self.amp_level,
-            )
+            model, optimizers = model.configure_apex(amp, model, self.optimizers, self.amp_level)
             self.optimizers = optimizers
 
         # DDP2 uses all GPUs on the machine
@@ -178,53 +175,12 @@ class TrainerDDPMixin(object):
         elif self.use_ddp2:
             device_ids = None
 
-        model = LightningDistributedDataParallel(
-            model,
-            device_ids=device_ids,
-            find_unused_parameters=True
-        )
+        # allow user to configure ddp
+        model = model.configure_ddp(model, device_ids)
 
         # continue training routine
         self.run_pretrain_routine(model)
 
-    def __init_tcp_connection(self):
-        """
-        Connect all procs in the world using the env:// init
-        Use the first node as the root address
-        :param port:
-        :param tries:
-        :return:
-        """
-
-        # use slurm job id for the port number
-        # guarantees unique ports across jobs from same grid search
-        try:
-            # use the last 4 numbers in the job id as the id
-            default_port = os.environ['SLURM_JOB_ID']
-            default_port = default_port[-4:]
-
-            # all ports should be in the 10k+ range
-            default_port = int(default_port) + 15000
-
-        except Exception as e:
-            default_port = 12910
-
-        # if user gave a port number, use that one instead
-        try:
-            default_port = os.environ['MASTER_PORT']
-        except Exception:
-            os.environ['MASTER_PORT'] = str(default_port)
-
-        # figure out the root node addr
-        try:
-            root_node = os.environ['SLURM_NODELIST'].split(' ')[0]
-        except Exception:
-            root_node = '127.0.0.2'
-
-        root_node = self.resolve_root_node_address(root_node)
-        os.environ['MASTER_ADDR'] = root_node
-        dist.init_process_group("nccl", rank=self.proc_rank, world_size=self.world_size)
-
     def resolve_root_node_address(self, root_node):
         if '[' in root_node:
             name = root_node.split('[')[0]
diff --git a/pytorch_lightning/trainer/dp_mixin.py b/pytorch_lightning/trainer/dp_mixin.py
index 0bde8a7b..684ff15c 100644
--- a/pytorch_lightning/trainer/dp_mixin.py
+++ b/pytorch_lightning/trainer/dp_mixin.py
@@ -71,9 +71,7 @@ class TrainerDPMixin(object):
 
         if self.use_amp:
             # An example
-            model, optimizers = amp.initialize(
-                model, self.optimizers, opt_level=self.amp_level,
-            )
+            model, optimizers = model.configure_apex(amp, model, self.optimizers, self.amp_level)
             self.optimizers = optimizers
 
         self.run_pretrain_routine(model)
diff --git a/pytorch_lightning/trainer/logging_mixin.py b/pytorch_lightning/trainer/logging_mixin.py
index 236e512d..13e5afb6 100644
--- a/pytorch_lightning/trainer/logging_mixin.py
+++ b/pytorch_lightning/trainer/logging_mixin.py
@@ -156,6 +156,10 @@ class TrainerLoggingMixin(object):
             if isinstance(output[k], dict):
                 output[k] = self.reduce_distributed_output(output[k], nb_gpus)
 
+            # do nothing when there's a scalar
+            elif isinstance(output[k], torch.Tensor) and output[k].dim() == 0:
+                pass
+
             # reduce only metrics that have the same nb of gpus
             elif output[k].size(0) == nb_gpus:
                 reduced = torch.mean(output[k])
diff --git a/pytorch_lightning/trainer/train_loop_mixin.py b/pytorch_lightning/trainer/train_loop_mixin.py
index d41fc7c1..306416db 100644
--- a/pytorch_lightning/trainer/train_loop_mixin.py
+++ b/pytorch_lightning/trainer/train_loop_mixin.py
@@ -189,13 +189,6 @@ class TrainerTrainLoopMixin(object):
                     callback_metrics = output[3]
                     self.hiddens = output[4]
 
-                    # track metrics for callbacks
-                    all_callback_metrics.append(callback_metrics)
-
-                    # track progress bar metrics
-                    self.add_tqdm_metrics(progress_bar_metrics)
-                    all_log_metrics.append(log_metrics)
-
                     # accumulate loss
                     # (if accumulate_grad_batches = 1 no effect)
                     closure_loss = closure_loss / self.accumulate_grad_batches
@@ -204,6 +197,13 @@ class TrainerTrainLoopMixin(object):
                     model_ref = self.get_model()
                     model_ref.backward(self.use_amp, closure_loss, optimizer)
 
+                    # track metrics for callbacks
+                    all_callback_metrics.append(callback_metrics)
+
+                    # track progress bar metrics
+                    self.add_tqdm_metrics(progress_bar_metrics)
+                    all_log_metrics.append(log_metrics)
+
                     # insert after step hook
                     if self.is_function_implemented('on_after_backward'):
                         model_ref = self.get_model()
@@ -277,13 +277,15 @@ class TrainerTrainLoopMixin(object):
         if len(self.optimizers) > 1:
             args.append(opt_idx)
 
+        # pass hiddens if using tbptt
         if self.truncated_bptt_steps is not None:
             args.append(hiddens)
 
-        if self.use_ddp or self.use_ddp2:
-            output = self.model(*args)
-        elif self.use_dp:
+        # distributed forward
+        if self.use_ddp or self.use_ddp2 or self.use_dp:
             output = self.model(*args)
+
+        # single GPU forward
         elif self.single_gpu:
             gpu_id = 0
             if type(self.data_parallel_device_ids) is list:
@@ -292,9 +294,16 @@ class TrainerTrainLoopMixin(object):
             args[0] = batch
             output = self.model.training_step(*args)
 
+        # CPU forward
         else:
             output = self.model.training_step(*args)
 
+        # allow any mode to define training_end
+        if self.is_overriden('training_end'):
+            model_ref = self.get_model()
+            output = model_ref.training_end(output)
+
         # format and reduce outputs accordingly
         output = self.process_output(output, train=True)
+
         return output

From 9fa28066059c3bda0b022c33921796c7425cd41e Mon Sep 17 00:00:00 2001
From: Nic Eggert <nic@eggert.io>
Date: Tue, 5 Nov 2019 09:41:59 -0600
Subject: [PATCH 5/6] Fix ModelCheckpoint default paths (#413)

* Make name and version properties required

* Warn before deleting files in checkpoint directory

* Get default checkpoint path from any logger

* Fix typos

* Uncomment logger tests

* Whitespace

* Update callback_config_mixin.py

checkpoints and version file names would just have a number. it's easy to tell what you're looking at with version_ prepended

* Address comments

* Fix broken tests
---
 pytorch_lightning/callbacks/pt_callbacks.py   |  10 +
 pytorch_lightning/logging/base.py             |   7 +-
 pytorch_lightning/logging/mlflow_logger.py    |   8 +
 pytorch_lightning/logging/test_tube_logger.py |  11 +-
 .../trainer/callback_config_mixin.py          |  15 +-
 tests/test_y_logging.py                       | 218 +++++++++---------
 6 files changed, 153 insertions(+), 116 deletions(-)

diff --git a/pytorch_lightning/callbacks/pt_callbacks.py b/pytorch_lightning/callbacks/pt_callbacks.py
index 2c1553d4..388fca69 100644
--- a/pytorch_lightning/callbacks/pt_callbacks.py
+++ b/pytorch_lightning/callbacks/pt_callbacks.py
@@ -179,6 +179,16 @@ class ModelCheckpoint(Callback):
                  save_best_only=True, save_weights_only=False,
                  mode='auto', period=1, prefix=''):
         super(ModelCheckpoint, self).__init__()
+        if (
+            save_best_only and
+            os.path.isdir(filepath) and
+            len(os.listdir(filepath)) > 0
+        ):
+            warnings.warn(
+                f"Checkpoint directory {filepath} exists and is not empty with save_best_only=True."
+                "All files in this directory will be deleted when a checkpoint is saved!"
+            )
+
         self.monitor = monitor
         self.verbose = verbose
         self.filepath = filepath
diff --git a/pytorch_lightning/logging/base.py b/pytorch_lightning/logging/base.py
index 351b945e..a72e072b 100644
--- a/pytorch_lightning/logging/base.py
+++ b/pytorch_lightning/logging/base.py
@@ -65,7 +65,12 @@ class LightningLoggerBase(object):
         """Set the process rank"""
         self._rank = value
 
+    @property
+    def name(self):
+        """Return the experiment name"""
+        raise NotImplementedError("Sub-classes must provide a name property")
+
     @property
     def version(self):
         """Return the experiment version"""
-        return None
+        raise NotImplementedError("Sub-classes must provide a version property")
diff --git a/pytorch_lightning/logging/mlflow_logger.py b/pytorch_lightning/logging/mlflow_logger.py
index 6c0d460a..e9b2abd5 100644
--- a/pytorch_lightning/logging/mlflow_logger.py
+++ b/pytorch_lightning/logging/mlflow_logger.py
@@ -60,3 +60,11 @@ class MLFlowLogger(LightningLoggerBase):
         if status == 'success':
             status = 'FINISHED'
         self.experiment.set_terminated(self.run_id, status)
+
+    @property
+    def name(self):
+        return self.experiment_name
+
+    @property
+    def version(self):
+        return self._run_id
diff --git a/pytorch_lightning/logging/test_tube_logger.py b/pytorch_lightning/logging/test_tube_logger.py
index 0d74307b..23e8806f 100644
--- a/pytorch_lightning/logging/test_tube_logger.py
+++ b/pytorch_lightning/logging/test_tube_logger.py
@@ -15,7 +15,7 @@ class TestTubeLogger(LightningLoggerBase):
     ):
         super().__init__()
         self.save_dir = save_dir
-        self.name = name
+        self._name = name
         self.description = description
         self.debug = debug
         self._version = version
@@ -29,7 +29,7 @@ class TestTubeLogger(LightningLoggerBase):
 
         self._experiment = Experiment(
             save_dir=self.save_dir,
-            name=self.name,
+            name=self._name,
             debug=self.debug,
             version=self.version,
             description=self.description,
@@ -80,6 +80,13 @@ class TestTubeLogger(LightningLoggerBase):
         if self._experiment is not None:
             self.experiment.rank = value
 
+    @property
+    def name(self):
+        if self._experiment is None:
+            return self._name
+        else:
+            return self.experiment.name
+
     @property
     def version(self):
         if self._experiment is None:
diff --git a/pytorch_lightning/trainer/callback_config_mixin.py b/pytorch_lightning/trainer/callback_config_mixin.py
index 1b36bd46..58e6ec05 100644
--- a/pytorch_lightning/trainer/callback_config_mixin.py
+++ b/pytorch_lightning/trainer/callback_config_mixin.py
@@ -1,3 +1,5 @@
+import os
+
 from pytorch_lightning.callbacks import ModelCheckpoint, EarlyStopping
 from pytorch_lightning.logging import TestTubeLogger
 
@@ -12,14 +14,15 @@ class TrainerCallbackConfigMixin(object):
         """
         if self.checkpoint_callback is True:
             # init a default one
-            if isinstance(self.logger, TestTubeLogger):
-                ckpt_path = '{}/{}/version_{}/{}'.format(
+            if self.logger is not None:
+                ckpt_path = os.path.join(
                     self.default_save_path,
-                    self.logger.experiment.name,
-                    self.logger.experiment.version,
-                    'checkpoints')
+                    self.logger.name,
+                    f'version_{self.logger.version}',
+                    "checkpoints"
+                )
             else:
-                ckpt_path = self.default_save_path
+                ckpt_path = os.path.join(self.default_save_path, "checkpoints")
 
             self.checkpoint_callback = ModelCheckpoint(
                 filepath=ckpt_path
diff --git a/tests/test_y_logging.py b/tests/test_y_logging.py
index 39b7edfc..c44a9e2c 100644
--- a/tests/test_y_logging.py
+++ b/tests/test_y_logging.py
@@ -1,3 +1,4 @@
+import os
 import pickle
 
 import numpy as np
@@ -5,6 +6,7 @@ import torch
 
 from pytorch_lightning import Trainer
 from pytorch_lightning.testing import LightningTestModel
+from pytorch_lightning.logging import LightningLoggerBase, rank_zero_only
 from . import testing_utils
 
 RANDOM_FILE_PATHS = list(np.random.randint(12000, 19000, 1000))
@@ -69,117 +71,119 @@ def test_testtube_pickle():
     testing_utils.clear_save_dir()
 
 
-# def test_mlflow_logger():
-#     """
-#     verify that basic functionality of mlflow logger works
-#     """
-#     reset_seed()
-#
-#     try:
-#         from pytorch_lightning.logging import MLFlowLogger
-#     except ModuleNotFoundError:
-#         return
-#
-#     hparams = get_hparams()
-#     model = LightningTestModel(hparams)
-#
-#     root_dir = os.path.dirname(os.path.realpath(__file__))
-#     mlflow_dir = os.path.join(root_dir, "mlruns")
-#     import pdb
-#     pdb.set_trace()
-#
-#     logger = MLFlowLogger("test", f"file://{mlflow_dir}")
-#     logger.log_hyperparams(hparams)
-#     logger.save()
-#
-#     trainer_options = dict(
-#         max_nb_epochs=1,
-#         train_percent_check=0.01,
-#         logger=logger
-#     )
-#
-#     trainer = Trainer(**trainer_options)
-#     result = trainer.fit(model)
-#
-#     print('result finished')
-#     assert result == 1, "Training failed"
-#
-#     shutil.move(mlflow_dir, mlflow_dir + f'_{n}')
+def test_mlflow_logger():
+    """
+    verify that basic functionality of mlflow logger works
+    """
+    reset_seed()
+
+    try:
+        from pytorch_lightning.logging import MLFlowLogger
+    except ModuleNotFoundError:
+        return
+
+    hparams = testing_utils.get_hparams()
+    model = LightningTestModel(hparams)
+
+    root_dir = os.path.dirname(os.path.realpath(__file__))
+    mlflow_dir = os.path.join(root_dir, "mlruns")
+
+    logger = MLFlowLogger("test", f"file://{mlflow_dir}")
+
+    trainer_options = dict(
+        max_nb_epochs=1,
+        train_percent_check=0.01,
+        logger=logger
+    )
+
+    trainer = Trainer(**trainer_options)
+    result = trainer.fit(model)
+
+    print('result finished')
+    assert result == 1, "Training failed"
+
+    testing_utils.clear_save_dir()
 
 
-# def test_mlflow_pickle():
-#     """
-#     verify that pickling trainer with mlflow logger works
-#     """
-#     reset_seed()
-#
-#     try:
-#         from pytorch_lightning.logging import MLFlowLogger
-#     except ModuleNotFoundError:
-#         return
-#
-#     hparams = get_hparams()
-#     model = LightningTestModel(hparams)
-#
-#     root_dir = os.path.dirname(os.path.realpath(__file__))
-#     mlflow_dir = os.path.join(root_dir, "mlruns")
-#
-#     logger = MLFlowLogger("test", f"file://{mlflow_dir}")
-#     logger.log_hyperparams(hparams)
-#     logger.save()
-#
-#     trainer_options = dict(
-#         max_nb_epochs=1,
-#         logger=logger
-#     )
-#
-#     trainer = Trainer(**trainer_options)
-#     pkl_bytes = pickle.dumps(trainer)
-#     trainer2 = pickle.loads(pkl_bytes)
-#     trainer2.logger.log_metrics({"acc": 1.0})
-#
-#     n = RANDOM_FILE_PATHS.pop()
-#     shutil.move(mlflow_dir, mlflow_dir + f'_{n}')
+def test_mlflow_pickle():
+    """
+    verify that pickling trainer with mlflow logger works
+    """
+    reset_seed()
+
+    try:
+        from pytorch_lightning.logging import MLFlowLogger
+    except ModuleNotFoundError:
+        return
+
+    hparams = testing_utils.get_hparams()
+    model = LightningTestModel(hparams)
+
+    root_dir = os.path.dirname(os.path.realpath(__file__))
+    mlflow_dir = os.path.join(root_dir, "mlruns")
+
+    logger = MLFlowLogger("test", f"file://{mlflow_dir}")
+
+    trainer_options = dict(
+        max_nb_epochs=1,
+        logger=logger
+    )
+
+    trainer = Trainer(**trainer_options)
+    pkl_bytes = pickle.dumps(trainer)
+    trainer2 = pickle.loads(pkl_bytes)
+    trainer2.logger.log_metrics({"acc": 1.0})
+
+    testing_utils.clear_save_dir()
 
 
-# def test_custom_logger():
-#
-#     class CustomLogger(LightningLoggerBase):
-#         def __init__(self):
-#             super().__init__()
-#             self.hparams_logged = None
-#             self.metrics_logged = None
-#             self.finalized = False
-#
-#         @rank_zero_only
-#         def log_hyperparams(self, params):
-#             self.hparams_logged = params
-#
-#         @rank_zero_only
-#         def log_metrics(self, metrics, step_num):
-#             self.metrics_logged = metrics
-#
-#         @rank_zero_only
-#         def finalize(self, status):
-#             self.finalized_status = status
-#
-#     hparams = get_hparams()
-#     model = LightningTestModel(hparams)
-#
-#     logger = CustomLogger()
-#
-#     trainer_options = dict(
-#         max_nb_epochs=1,
-#         train_percent_check=0.01,
-#         logger=logger
-#     )
-#
-#     trainer = Trainer(**trainer_options)
-#     result = trainer.fit(model)
-#     assert result == 1, "Training failed"
-#     assert logger.hparams_logged == hparams
-#     assert logger.metrics_logged != {}
-#     assert logger.finalized_status == "success"
+def test_custom_logger(tmpdir):
+
+    class CustomLogger(LightningLoggerBase):
+        def __init__(self):
+            super().__init__()
+            self.hparams_logged = None
+            self.metrics_logged = None
+            self.finalized = False
+
+        @rank_zero_only
+        def log_hyperparams(self, params):
+            self.hparams_logged = params
+
+        @rank_zero_only
+        def log_metrics(self, metrics, step_num):
+            self.metrics_logged = metrics
+
+        @rank_zero_only
+        def finalize(self, status):
+            self.finalized_status = status
+
+        @property
+        def name(self):
+            return "name"
+
+        @property
+        def version(self):
+            return "1"
+
+    hparams = testing_utils.get_hparams()
+    model = LightningTestModel(hparams)
+
+    logger = CustomLogger()
+
+    trainer_options = dict(
+        max_nb_epochs=1,
+        train_percent_check=0.01,
+        logger=logger,
+        default_save_path=tmpdir
+    )
+
+    trainer = Trainer(**trainer_options)
+    result = trainer.fit(model)
+    assert result == 1, "Training failed"
+    assert logger.hparams_logged == hparams
+    assert logger.metrics_logged != {}
+    assert logger.finalized_status == "success"
 
 
 def reset_seed():

From f9a3e03e1c89a5085366454206fef2e5459c4cff Mon Sep 17 00:00:00 2001
From: William Falcon <waf2107@columbia.edu>
Date: Tue, 5 Nov 2019 10:49:22 -0500
Subject: [PATCH 6/6] Update CONTRIBUTING.md

---
 .github/CONTRIBUTING.md | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/.github/CONTRIBUTING.md b/.github/CONTRIBUTING.md
index 109c5b7c..e9182fb2 100644
--- a/.github/CONTRIBUTING.md
+++ b/.github/CONTRIBUTING.md
@@ -30,6 +30,9 @@ You shouldn't be afraid to upgrade Lightning :)
 #### Gain User Trust    
 As a researcher you can't have any part of your code going wrong. So, make thorough tests that ensure an implementation of a new trick or subbtle change is correct.    
 
+#### Interoperability  
+Have a favorite feature from other libraries like fast.ai or transformers? Those should just work with lightning as well. Grab your favorite model or learning rate scheduler from your favorite library and run it in Lightning.
+
 ## Contribution Types    
 Currently looking for help implementing new features or adding bug fixes.