bug free release

* fixed ckpt tests

* fixed ckpt tests

* fixed ckpt tests

* fixed ckpt tests

* fixed ckpt tests

* fixed ckpt tests

* fixed ckpt tests

* fixed ckpt tests

* fixed ckpt tests

* fixed ckpt tests

* fixed ckpt tests

* fixed ckpt tests

* fixed ckpt tests

* fixed ckpt tests

* fixed ckpt tests

* fixed ckpt tests

* fixed ckpt tests

* fixed ckpt tests

* fixed ckpt tests

* fixed ckpt tests

* fixed ckpt tests

* fixed ckpt tests

* fixed ckpt tests

* fixed ckpt tests

* fixed ckpt tests

* fixed ckpt tests

* fixed ckpt tests

* fixed ckpt tests

* fixed ckpt tests

* fixed ckpt tests

* fixed ckpt tests

* fixed ckpt tests

* fixed ckpt tests

* fixed ckpt tests

* fixed ckpt tests

* fixed ckpt tests

* fixed ckpt tests

* fixed ckpt tests

* fixed ckpt tests

* fixed ckpt tests

* fixed ckpt tests

* fixed ckpt tests

* fixed ckpt tests

* fixed ckpt tests

* fixed ckpt tests

* fixed ckpt tests

* fixed ckpt tests

* fixed ckpt tests

* fixed ckpt tests

* fixed ckpt tests

* fixed ckpt tests

* fixed ckpt tests

* fixed ckpt tests

* fixed ckpt tests

* fixed ckpt tests

* fixed ckpt tests

* fixed ckpt tests

* fixed ckpt tests

* fixed ckpt tests

* fixed ckpt tests

* fixed ckpt tests

* fixed ckpt tests

* fixed ckpt tests

* fixed ckpt tests

* fixed ckpt tests

README.md

@@ -17,6 +17,8 @@
 [![ReadTheDocs](https://readthedocs.org/projects/pytorch-lightning/badge/?version=latest)](https://pytorch-lightning.readthedocs.io/en/latest)
 [![Gitter](https://badges.gitter.im/PyTorch-Lightning/community.svg)](https://gitter.im/PyTorch-Lightning/community?utm_source=badge&utm_medium=badge&utm_campaign=pr-badge)
 [![license](https://img.shields.io/badge/License-Apache%202.0-blue.svg)](https://github.com/williamFalcon/pytorch-lightning/blob/master/LICENSE)
+[![Next Release](https://img.shields.io/badge/Next%20Release-Nov%206-<COLOR>.svg)](https://shields.io/)
+
 <!-- 
 removed until codecov badge isn't empy. likely a config error showing nothing on master.
 [![codecov](https://codecov.io/gh/Borda/pytorch-lightning/branch/master/graph/badge.svg)](https://codecov.io/gh/Borda/pytorch-lightning)
@@ -52,7 +54,7 @@ Lightning sets up all the boilerplate state-of-the-art training for you so you c
 - [Lightning features](https://github.com/williamFalcon/pytorch-lightning#lightning-automates-all-of-the-following-each-is-also-configurable)    
 - [Examples](https://github.com/williamFalcon/pytorch-lightning#examples)    
 - [Tutorials](https://github.com/williamFalcon/pytorch-lightning#tutorials)
-- [Contributing](https://github.com/williamFalcon/pytorch-lightning/blob/master/CONTRIBUTING.md)    
+- [Contributing](https://github.com/williamFalcon/pytorch-lightning/blob/master/.github/CONTRIBUTING.md)
 - [Bleeding edge install](https://github.com/williamFalcon/pytorch-lightning#bleeding-edge)   
 - [Lightning Design Principles](https://github.com/williamFalcon/pytorch-lightning#lightning-design-principles)   
 - [Asking for help](https://github.com/williamFalcon/pytorch-lightning#asking-for-help)
@@ -265,24 +267,8 @@ Lightning also adds a text column with all the hyperparameters for this experime
 
 ![tensorboard-support](./docs/source/_static/tf_tags.png)
 
-Simply note the path you set for the [Experiment](https://williamfalcon.github.io/test-tube/experiment_tracking/experiment/) from [test_tube](https://github.com/williamFalcon/test-tube)
-```python
-from test_tube import Experiment
-from pytorch_lightning import Trainer
-
-exp = Experiment(save_dir='/some/path')
-trainer = Trainer(experiment=exp)
-...
-```   
-
-And run tensorboard from that dir   
-```bash
-tensorboard --logdir /some/path     
-```    
-
 ## Lightning automates all of the following ([each is also configurable](https://williamfalcon.github.io/pytorch-lightning/Trainer/)):
 
-
 #### Checkpointing    
 
 - [Checkpoint callback](https://williamfalcon.github.io/pytorch-lightning/Trainer/Checkpointing/#model-saving)

docs/LightningModule/RequiredTrainerInterface.md

@@ -77,7 +77,7 @@ class CoolModel(pl.LightningModule):
 
     def configure_optimizers(self):
         # REQUIRED
-        return [torch.optim.Adam(self.parameters(), lr=0.02)]
+        return torch.optim.Adam(self.parameters(), lr=0.02)
 
     @pl.data_loader
     def train_dataloader(self):

docs/Trainer/debugging.md

@@ -38,6 +38,14 @@ trainer = Trainer(overfit_pct=0.01)
 #### Print the parameter count by layer
 By default lightning prints a list of parameters *and submodules* when it starts training.
 
+``` {.python}
+# DEFAULT print a full list of all submodules and their parameters.
+trainer = Trainer(weights_summary='full')
+
+# only print the top-level modules (i.e. the children of LightningModule).
+trainer = Trainer(weights_summary='top')
+```
+
 ---
 #### Print which gradients are nan 
 This option prints a list of tensors with nan gradients.

docs/Trainer/hooks.md

@@ -65,12 +65,12 @@ You can override this method to adjust how you do the optimizer step for each op
 Called once per optimizer
 ```python
 # DEFAULT
-def optimizer_step(self, current_epoch, batch_nb, optimizer, optimizer_i):
+def optimizer_step(self, current_epoch, batch_nb, optimizer, optimizer_i, second_order_closure=None):
     optimizer.step()   
     optimizer.zero_grad()   
     
 # Alternating schedule for optimizer steps (ie: GANs)    
-def optimizer_step(self, current_epoch, batch_nb, optimizer, optimizer_i):
+def optimizer_step(self, current_epoch, batch_nb, optimizer, optimizer_i, second_order_closure=None):
     # update generator opt every 2 steps
     if optimizer_i == 0:
         if batch_nb % 2 == 0 :
@@ -91,7 +91,7 @@ This step allows you to do a lot of non-standard training tricks such as learnin
 
 ```python
 # learning rate warm-up
-def optimizer_step(self, current_epoch, batch_nb, optimizer, optimizer_i):
+def optimizer_step(self, current_epoch, batch_nb, optimizer, optimizer_i, second_order_closure=None):
     # warm up lr
     if self.trainer.global_step < 500:
         lr_scale = min(1., float(self.trainer.global_step + 1) / 500.)

docs/examples/Examples.md

@@ -1,5 +1,5 @@
 ### Template model definition
-In 99% of cases you want to just copy [this template](https://github.com/williamFalcon/pytorch-lightning/blob/master/examples/new_project_templates/lightning_module_template.py) to start a new lightningModule and change the core of what your model is actually trying to do.
+In 99% of cases you want to just copy [one of the examples](https://github.com/williamFalcon/pytorch-lightning/tree/master/examples) to start a new lightningModule and change the core of what your model is actually trying to do.
 
 ```bash
 # get a copy of the module template

docs/index.md

@@ -60,9 +60,8 @@ Notice a few things about this flow:
 ###### Templates 
 1. [MNIST LightningModule](https://williamfalcon.github.io/pytorch-lightning/LightningModule/RequiredTrainerInterface/#minimal-example) 
 2. [Trainer](https://williamfalcon.github.io/pytorch-lightning/Trainer/)
-    - [Basic CPU Trainer Template](https://github.com/williamFalcon/pytorch-lightning/blob/master/examples/new_project_templates/single_cpu_template.py) 
-    - [Multi-GPU Trainer Template](https://github.com/williamFalcon/pytorch-lightning/blob/master/examples/new_project_templates/single_gpu_node_template.py)
-    - [GPU cluster Trainer Template](https://github.com/williamFalcon/pytorch-lightning/blob/master/examples/new_project_templates/multi_node_cluster_template.py)
+    - [Basic CPU, GPU Trainer Template](https://github.com/williamFalcon/pytorch-lightning/tree/master/examples/basic_examples) 
+    - [GPU cluster Trainer Template](https://github.com/williamFalcon/pytorch-lightning/tree/master/examples/multi_node_examples)
 
 ###### Docs shortcuts
 - [LightningModule](LightningModule/RequiredTrainerInterface/)  

examples/basic_examples/lightning_module_template.py

@@ -98,8 +98,11 @@ class LightningTemplateModel(LightningModule):
         if self.trainer.use_dp or self.trainer.use_ddp2:
             loss_val = loss_val.unsqueeze(0)
 
+        tqdm_dict = {'train_loss': loss_val}
         output = OrderedDict({
-            'loss': loss_val
+            'loss': loss_val,
+            'progress_bar': tqdm_dict,
+            'log': tqdm_dict
         })
 
         # can also return just a scalar instead of a dict (return loss_val)
@@ -168,7 +171,7 @@ class LightningTemplateModel(LightningModule):
         val_loss_mean /= len(outputs)
         val_acc_mean /= len(outputs)
         tqdm_dict = {'val_loss': val_loss_mean, 'val_acc': val_acc_mean}
-        result = {'progress_bar': tqdm_dict, 'logs': tqdm_dict}
+        result = {'progress_bar': tqdm_dict, 'log': tqdm_dict}
         return result
 
     # ---------------------

examples/multi_node_examples/README.md

@@ -1,6 +1,7 @@
 # Multi-node example   
 
-To run this demo which launches a single job that trains on 2 nodes (2 gpus per node), do the following:
+This demo launches a job using 2 GPUs on 2 different nodes (4 GPUs total).
+To run this demo do the following:
 
 1. Log into the jumphost node of your SLURM-managed cluster.  
 2. Create a conda environment with Lightning and a GPU PyTorch version.   

pytorch_lightning/callbacks/pt_callbacks.py

@@ -125,7 +125,8 @@ class EarlyStopping(Callback):
             print('Early stopping conditioned on metric `%s` '
                   'which is not available. Available metrics are: %s' %
                   (self.monitor, ','.join(list(logs.keys()))), RuntimeWarning)
-            exit(-1)
+            stop_training = True
+            return stop_training
 
         if self.monitor_op(current - self.min_delta, self.best):
             self.best = current

pytorch_lightning/logging/mlflow_logger.py

@@ -9,11 +9,12 @@ logger = getLogger(__name__)
 
 
 class MLFlowLogger(LightningLoggerBase):
-    def __init__(self, experiment_name, tracking_uri=None):
+    def __init__(self, experiment_name, tracking_uri=None, tags=None):
         super().__init__()
         self.client = mlflow.tracking.MlflowClient(tracking_uri)
         self.experiment_name = experiment_name
         self._run_id = None
+        self.tags = tags
 
     @property
     def run_id(self):
@@ -28,7 +29,7 @@ class MLFlowLogger(LightningLoggerBase):
             self.client.create_experiment(self.experiment_name)
             experiment = self.client.get_experiment_by_name(self.experiment_name)
 
-        run = self.client.create_run(experiment.experiment_id)
+        run = self.client.create_run(experiment.experiment_id, tags=self.tags)
         self._run_id = run.info.run_id
         return self._run_id
 
@@ -53,4 +54,6 @@ class MLFlowLogger(LightningLoggerBase):
 
     @rank_zero_only
     def finalize(self, status="FINISHED"):
+        if status == 'success':
+            status = 'FINISHED'
         self.client.set_terminated(self.run_id, status)

pytorch_lightning/logging/test_tube_logger.py

@@ -10,11 +10,13 @@ class TestTubeLogger(LightningLoggerBase):
     __test__ = False
 
     def __init__(
-        self, save_dir, name="default", debug=False, version=None, create_git_tag=False
+        self, save_dir, name="default", description=None, debug=False,
+            version=None, create_git_tag=False
     ):
         super().__init__()
         self.save_dir = save_dir
         self.name = name
+        self.description = description
         self.debug = debug
         self._version = version
         self.create_git_tag = create_git_tag
@@ -24,11 +26,13 @@ class TestTubeLogger(LightningLoggerBase):
     def experiment(self):
         if self._experiment is not None:
             return self._experiment
+
         self._experiment = Experiment(
             save_dir=self.save_dir,
             name=self.name,
             debug=self.debug,
             version=self.version,
+            description=self.description,
             create_git_tag=self.create_git_tag,
             rank=self.rank,
         )
@@ -36,37 +40,45 @@ class TestTubeLogger(LightningLoggerBase):
 
     @rank_zero_only
     def log_hyperparams(self, params):
+        # TODO: HACK figure out where this is being set to true
+        self.experiment.debug = self.debug
         self.experiment.argparse(params)
 
     @rank_zero_only
     def log_metrics(self, metrics, step_num=None):
+        # TODO: HACK figure out where this is being set to true
+        self.experiment.debug = self.debug
         self.experiment.log(metrics, global_step=step_num)
 
     @rank_zero_only
     def save(self):
+        # TODO: HACK figure out where this is being set to true
+        self.experiment.debug = self.debug
         self.experiment.save()
 
     @rank_zero_only
     def finalize(self, status):
+        # TODO: HACK figure out where this is being set to true
+        self.experiment.debug = self.debug
         self.save()
         self.close()
 
+    @rank_zero_only
     def close(self):
-        self.experiment.close()
+        # TODO: HACK figure out where this is being set to true
+        self.experiment.debug = self.debug
+        exp = self.experiment
+        exp.close()
 
     @property
     def rank(self):
-        if self._experiment is None:
-            return self._rank
-        else:
-            return self.experiment.rank
+        return self._rank
 
     @rank.setter
     def rank(self, value):
-        if self._experiment is None:
-            self._rank = value
-        else:
-            return self.experiment.rank
+        self._rank = value
+        if self._experiment is not None:
+            self.experiment.rank = value
 
     @property
     def version(self):

pytorch_lightning/root_module/memory.py

@@ -12,11 +12,12 @@ import pandas as pd
 
 class ModelSummary(object):
 
-    def __init__(self, model):
+    def __init__(self, model, mode='full'):
         '''
         Generates summaries of model layers and dimensions.
         '''
         self.model = model
+        self.mode = mode
         self.in_sizes = []
         self.out_sizes = []
 
@@ -28,9 +29,20 @@ class ModelSummary(object):
     def __repr__(self):
         return self.summary.__str__()
 
+    def named_modules(self):
+        if self.mode == 'full':
+            mods = self.model.named_modules()
+            mods = list(mods)[1:]  # do not include root module (LightningModule)
+        elif self.mode == 'top':
+            # the children are the top-level modules
+            mods = self.model.named_children()
+        else:
+            mods = []
+        return list(mods)
+
     def get_variable_sizes(self):
         '''Run sample input through each layer to get output sizes'''
-        mods = list(self.model.modules())
+        mods = self.named_modules()
         in_sizes = []
         out_sizes = []
         input_ = self.model.example_input_array
@@ -43,8 +55,7 @@ class ModelSummary(object):
 
         with torch.no_grad():
 
-            for i in range(1, len(mods)):
-                m = mods[i]
+            for _, m in mods:
                 if type(input_) is list or type(input_) is tuple:  # pragma: no cover
                     out = m(*input_)
                 else:
@@ -72,16 +83,17 @@ class ModelSummary(object):
 
         self.in_sizes = in_sizes
         self.out_sizes = out_sizes
+        assert len(in_sizes) == len(out_sizes)
         return
 
     def get_layer_names(self):
         '''Collect Layer Names'''
-        mods = list(self.model.named_modules())
+        mods = self.named_modules()
         names = []
         layers = []
-        for m in mods[1:]:
-            names += [m[0]]
-            layers += [str(m[1].__class__)]
+        for name, m in mods:
+            names += [name]
+            layers += [str(m.__class__)]
 
         layer_types = [x.split('.')[-1][:-2] for x in layers]
 
@@ -91,11 +103,9 @@ class ModelSummary(object):
 
     def get_parameter_sizes(self):
         '''Get sizes of all parameters in `model`'''
-        mods = list(self.model.modules())
+        mods = self.named_modules()
         sizes = []
-
-        for i in range(1, len(mods)):
-            m = mods[i]
+        for _, m in mods:
             p = list(m.parameters())
             modsz = []
             for j in range(len(p)):
@@ -133,6 +143,7 @@ class ModelSummary(object):
         df['Name'] = self.layer_names
         df['Type'] = self.layer_types
         df['Params'] = self.param_nums
+        df['Params'] = df['Params'].map(get_human_readable_count)
 
         if self.model.example_input_array is not None:
 
@@ -226,3 +237,28 @@ def get_gpu_memory_map():
         k = f'gpu_{k}'
         gpu_memory_map[k] = v
     return gpu_memory_map
+
+
+def get_human_readable_count(number):
+    """
+    Abbreviates an integer number with K, M, B, T for thousands, millions,
+    billions and trillions, respectively.
+    Examples:
+        123     -> 123
+        1234    -> 1 K       (one thousand)
+        2e6     -> 2 M       (two million)
+        3e9     -> 3 B       (three billion)
+        4e12    -> 4 T       (four trillion)
+        5e15    -> 5,000 T
+    :param number: a positive integer number
+    :returns a string formatted according to the pattern described above.
+    """
+    assert number >= 0
+    labels = [' ', 'K', 'M', 'B', 'T']
+    num_digits = int(np.floor(np.log10(number)) + 1 if number > 0 else 1)
+    num_groups = int(np.ceil(num_digits / 3))
+    num_groups = min(num_groups, len(labels))  # don't abbreviate beyond trillions
+    shift = -3 * (num_groups - 1)
+    number = number * (10 ** shift)
+    index = num_groups - 1
+    return f'{int(number):,d} {labels[index]}'

pytorch_lightning/root_module/root_module.py

@@ -159,8 +159,8 @@ class LightningModule(GradInformation, ModelIO, ModelHooks):
 
         return model
 
-    def summarize(self):
-        model_summary = ModelSummary(self)
+    def summarize(self, mode):
+        model_summary = ModelSummary(self, mode=mode)
         print(model_summary)
 
     def freeze(self):

pytorch_lightning/testing/lm_test_module_base.py

@@ -104,7 +104,8 @@ class LightningTestModelBase(LightningModule):
         if self.trainer.batch_nb % 1 == 0:
             output = OrderedDict({
                 'loss': loss_val,
-                'progress_bar': {'some_val': loss_val * loss_val}
+                'progress_bar': {'some_val': loss_val * loss_val},
+                'log': {'train_some_val': loss_val * loss_val},
             })
 
             return output

pytorch_lightning/testing/lm_test_module_mixins.py

@@ -105,7 +105,7 @@ class LightningValidationMixin(LightningValidationStepMixin):
         val_acc_mean /= len(outputs)
 
         tqdm_dict = {'val_loss': val_loss_mean.item(), 'val_acc': val_acc_mean.item()}
-        results = {'progress_bar': tqdm_dict}
+        results = {'progress_bar': tqdm_dict, 'log': tqdm_dict}
         return results
 
 

pytorch_lightning/trainer/trainer.py

@@ -84,9 +84,9 @@ class Trainer(TrainerIO):
                  distributed_backend=None,
                  use_amp=False,
                  print_nan_grads=False,
-                 print_weights_summary=True,
+                 weights_summary='full',
                  weights_save_path=None,
-                 amp_level='O2',
+                 amp_level='O1',
                  nb_sanity_val_steps=5):
         """
 
@@ -116,7 +116,7 @@ class Trainer(TrainerIO):
         :param distributed_backend: str. Options: 'dp', 'ddp', 'ddp2'.
         :param use_amp: Bool. If true uses apex for 16bit precision
         :param print_nan_grads: Bool. Prints nan gradients
-        :param print_weights_summary: Bool. Prints summary of weights
+        :param weights_summary: str. Options: 'full', 'top', None to not print.
         :param weights_save_path: Bool. Where to save weights if on cluster
         :param amp_level: str. Check nvidia docs for level
         :param nb_sanity_val_steps: int. How many val steps before a full train loop.
@@ -128,15 +128,24 @@ class Trainer(TrainerIO):
         self.check_val_every_n_epoch = check_val_every_n_epoch
         self.enable_early_stop = early_stop_callback is not None
         self.track_grad_norm = track_grad_norm
-        self.fast_dev_run = fast_dev_run
         self.on_gpu = gpus is not None and torch.cuda.is_available()
         self.process_position = process_position
-        self.print_weights_summary = print_weights_summary
+        self.weights_summary = weights_summary
         self.max_nb_epochs = max_nb_epochs
         self.min_nb_epochs = min_nb_epochs
         self.nb_sanity_val_steps = nb_sanity_val_steps
         self.print_nan_grads = print_nan_grads
 
+        self.fast_dev_run = fast_dev_run
+        if self.fast_dev_run:
+            self.nb_sanity_val_steps = 1
+            self.max_nb_epochs = 1
+            m = '''
+            Running in fast_dev_run mode: will run a full train,
+            val loop using a single batch
+            '''
+            print(m)
+
         # set default save path if user didn't provide one
         self.default_save_path = default_save_path
         if self.default_save_path is None:
@@ -148,6 +157,7 @@ class Trainer(TrainerIO):
         self.avg_loss = 0
         self.batch_nb = 0
         self.tqdm_metrics = {}
+        self.callback_metrics = {}
         self.nb_val_batches = 0
         self.nb_training_batches = 0
         self.nb_test_batches = 0
@@ -183,22 +193,12 @@ class Trainer(TrainerIO):
                 version=self.slurm_job_id,
                 name='lightning_logs'
             )
+        self.logger.rank = 0
 
         # configure checkpoint callback
         self.checkpoint_callback = checkpoint_callback
-        if self.checkpoint_callback is None:
-            if isinstance(logger, TestTubeLogger):
-                ckpt_path = '{}/{}/{}'.format(self.default_save_path, self.logger.name,
-                                              self.logger.version)
-            else:
-                ckpt_path = self.default_save_path
 
-            self.checkpoint_callback = ModelCheckpoint(
-                filepath=ckpt_path
-            )
-
-        # configure weights save path
-        self.__configure_weights_path(checkpoint_callback, weights_save_path)
+        self.weights_save_path = weights_save_path
 
         # accumulated grads
         self.__configure_accumulated_gradients(accumulate_grad_batches)
@@ -250,22 +250,35 @@ class Trainer(TrainerIO):
             job_id = None
         return job_id
 
-    def __configure_weights_path(self, checkpoint_callback, weights_save_path):
+    def __configure_checkpoint_callback(self):
         """
         Weight path set in this priority:
         Checkpoint_callback's path (if passed in).
         User provided weights_saved_path
         Otherwise use os.getcwd()
         """
-        self.weights_save_path = weights_save_path
+        if self.checkpoint_callback is None:
+            # init a default one
+            if isinstance(self.logger, TestTubeLogger):
+                ckpt_path = '{}/{}/version_{}/{}'.format(
+                    self.default_save_path,
+                    self.logger.experiment.name,
+                    self.logger.experiment.version,
+                    'checkpoints')
+            else:
+                ckpt_path = self.default_save_path
 
-        if self.checkpoint_callback is not None:
-            self.checkpoint_callback.save_function = self.save_checkpoint
+            self.checkpoint_callback = ModelCheckpoint(
+                filepath=ckpt_path
+            )
 
-            # if checkpoint callback used, then override the weights path
-            self.weights_save_path = self.checkpoint_callback.filepath
+        # set the path for the callbacks
+        self.checkpoint_callback.save_function = self.save_checkpoint
 
-        # if weights_save_path is still none here, set to current workingdir
+        # if checkpoint callback used, then override the weights path
+        self.weights_save_path = self.checkpoint_callback.filepath
+
+        # if weights_save_path is still none here, set to current working dir
         if self.weights_save_path is None:
             self.weights_save_path = self.default_save_path
 
@@ -659,11 +672,12 @@ class Trainer(TrainerIO):
                 """
                 warnings.warn(msg)
 
-            if on_ddp and self.get_val_dataloaders is not None:
+            if on_ddp and self.get_val_dataloaders() is not None:
                 for dataloader in self.get_val_dataloaders():
                     if not isinstance(dataloader.sampler, DistributedSampler):
                         msg = """
                         Your val_dataloader(s) don't use DistributedSampler.
+
                         You're using multiple gpus and multiple nodes without using a
                         DistributedSampler to assign a subset of your data to each process.
                         To silence this warning, pass a DistributedSampler to your DataLoader.
@@ -682,11 +696,12 @@ class Trainer(TrainerIO):
                         warnings.warn(msg)
                         break
 
-            if on_ddp and self.get_test_dataloaders is not None:
+            if on_ddp and self.get_test_dataloaders() is not None:
                 for dataloader in self.get_test_dataloaders():
                     if not isinstance(dataloader.sampler, DistributedSampler):
                         msg = """
                         Your test_dataloader(s) don't use DistributedSampler.
+
                         You're using multiple gpus and multiple nodes without using a
                         DistributedSampler to assign a subset of your data to each process.
                         To silence this warning, pass a DistributedSampler to your DataLoader.
@@ -968,6 +983,22 @@ class Trainer(TrainerIO):
         ref_model.use_amp = self.use_amp
         ref_model.testing = self.testing
 
+        # link up experiment object
+        if self.logger is not None:
+            ref_model.logger = self.logger
+
+            # save exp to get started
+            if hasattr(ref_model, "hparams"):
+                self.logger.log_hyperparams(ref_model.hparams)
+
+            self.logger.save()
+
+        if self.use_ddp or self.use_ddp2:
+            dist.barrier()
+
+        # set up checkpoint callback
+        self.__configure_checkpoint_callback()
+
         # register auto-resubmit when on SLURM
         self.register_slurm_signal_handlers()
 
@@ -978,17 +1009,12 @@ class Trainer(TrainerIO):
         self.__layout_bookeeping()
 
         # print model summary
-        if self.proc_rank == 0 and self.print_weights_summary:
-            ref_model.summarize()
-
-        # link up experiment object
-        if self.logger is not None:
-            ref_model.logger = self.logger
-
-            # save exp to get started
-            if hasattr(ref_model, "hparams"):
-                self.logger.log_hyperparams(ref_model.hparams)
-            self.logger.save()
+        if self.proc_rank == 0 and self.weights_summary is not None:
+            if self.weights_summary in ['full', 'top']:
+                ref_model.summarize(mode=self.weights_summary)
+            else:
+                m = "weights_summary can be None, 'full' or 'top'"
+                raise MisconfigurationException(m)
 
         # track model now.
         # if cluster resets state, the model will update with the saved weights
@@ -1037,6 +1063,9 @@ class Trainer(TrainerIO):
             self.total_batches = self.nb_training_batches + self.nb_val_batches
             self.batch_loss_value = 0  # accumulated grads
 
+            # limit the number of batches to 1 in fast_dev_run
+            self.total_batches = 1
+
             # init progress_bar when requested
             if self.show_progress_bar:
                 self.progress_bar.reset(self.total_batches)
@@ -1056,14 +1085,17 @@ class Trainer(TrainerIO):
 
             # early stopping
             met_min_epochs = epoch_nb > self.min_nb_epochs
-            if self.enable_early_stop and met_min_epochs:
+            if self.enable_early_stop and (met_min_epochs or self.fast_dev_run):
                 should_stop = self.early_stop_callback.on_epoch_end(epoch=epoch_nb,
-                                                                    logs=self.__training_tqdm_dict)
+                                                                    logs=self.callback_metrics)
                 # stop training
                 stop = should_stop and met_min_epochs
                 if stop:
                     return
 
+        if self.logger is not None:
+            self.logger.finalize("success")
+
     def run_training_epoch(self):
         # before epoch hook
         if self.__is_function_implemented('on_epoch_start'):
@@ -1097,23 +1129,27 @@ class Trainer(TrainerIO):
             # ---------------
             is_val_check_batch = (batch_nb + 1) % self.val_check_batch == 0
             can_check_epoch = (self.current_epoch + 1) % self.check_val_every_n_epoch == 0
-            if self.fast_dev_run or is_val_check_batch or early_stop_epoch:
-                if can_check_epoch:
-                    self.__run_evaluation(test=self.testing)
+            should_check_val = ((is_val_check_batch or early_stop_epoch) and can_check_epoch)
 
-            # when batch should be saved
-            if (batch_nb + 1) % self.log_save_interval == 0 or early_stop_epoch:
+            # fast_dev_run always forces val checking after train batch
+            if self.fast_dev_run or should_check_val:
+                self.__run_evaluation(test=self.testing)
+
+            # when logs should be saved
+            should_save_log = (batch_nb + 1) % self.log_save_interval == 0 or early_stop_epoch
+            if should_save_log or self.fast_dev_run:
                 if self.proc_rank == 0 and self.logger is not None:
                     self.logger.save()
 
             # when metrics should be logged
-            if batch_nb % self.row_log_interval == 0 or early_stop_epoch:
+            should_log_metrics = batch_nb % self.row_log_interval == 0 or early_stop_epoch
+            if should_log_metrics or self.fast_dev_run:
 
                 # logs user requested information to logger
                 self.__log_metrics(batch_step_metrics, grad_norm_dic)
 
             # end epoch early
-            if early_stop_epoch:
+            if early_stop_epoch or self.fast_dev_run:
                 break
 
         # epoch end hook
@@ -1157,12 +1193,14 @@ class Trainer(TrainerIO):
     def __metrics_to_scalars(self, metrics):
         new_metrics = {}
         for k, v in metrics.items():
-            if type(v) is torch.Tensor:
+            if isinstance(v, torch.Tensor):
                 v = v.item()
 
             if type(v) is dict:
                 v = self.__metrics_to_scalars(v)
 
+            new_metrics[k] = v
+
         return new_metrics
 
     def __log_vals_blacklist(self):
@@ -1232,8 +1270,9 @@ class Trainer(TrainerIO):
             output = self.model.training_step(*args)
 
         # format and reduce outputs accordingly
-        loss, progress_bar_metrics, log_metrics = self.__process_output(output, train=True)
-        return loss, progress_bar_metrics, log_metrics
+        output = self.__process_output(output, train=True)
+        loss, progress_bar_metrics, log_metrics, callback_metrics = output
+        return loss, progress_bar_metrics, log_metrics, callback_metrics
 
     def __process_output(self, output, train=False):
         """
@@ -1242,6 +1281,25 @@ class Trainer(TrainerIO):
         :param output:
         :return:
         """
+        # ---------------
+        # EXTRACT CALLBACK KEYS
+        # ---------------
+        # all keys not progress_bar or log are candidates for callbacks
+        callback_metrics = {}
+        for k, v in output.items():
+            if k not in ['progress_bar', 'log']:
+                callback_metrics[k] = v
+
+        if train and self.use_dp or self.use_ddp2:
+            nb_gpus = self.num_gpus
+            callback_metrics = reduce_distributed_output(callback_metrics, nb_gpus)
+
+        for k, v in callback_metrics.items():
+            callback_metrics[k] = v.item()
+
+        # ---------------
+        # EXTRACT PROGRESS BAR KEYS
+        # ---------------
         try:
             progress_output = output['progress_bar']
 
@@ -1254,6 +1312,9 @@ class Trainer(TrainerIO):
         except Exception:
             progress_bar_metrics = {}
 
+        # ---------------
+        # EXTRACT LOGGING KEYS
+        # ---------------
         # extract metrics to log to experiment
         try:
             log_output = output['log']
@@ -1288,7 +1349,7 @@ class Trainer(TrainerIO):
             if self.use_dp or self.use_ddp2:
                 loss = reduce_distributed_output(loss, self.num_gpus)
 
-        return loss, progress_bar_metrics, log_metrics
+        return loss, progress_bar_metrics, log_metrics, callback_metrics
 
     def __clip_gradients(self):
         if self.gradient_clip_val > 0:
@@ -1305,6 +1366,9 @@ class Trainer(TrainerIO):
         # track grad norms
         grad_norm_dic = {}
 
+        # track all metrics for callbacks
+        all_callback_metrics = []
+
         # track metrics to log
         all_log_metrics = []
 
@@ -1329,11 +1393,13 @@ class Trainer(TrainerIO):
             def optimizer_closure():
                 # forward pass
                 output = self.__training_forward(batch, batch_nb, opt_idx)
-                closure_loss, progress_bar_metrics, log_metrics = output
+                closure_loss, progress_bar_metrics, log_metrics, callback_metrics = output
+
+                # 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
@@ -1401,6 +1467,9 @@ class Trainer(TrainerIO):
         # collapse all metrics into one dict
         all_log_metrics = {k: v for d in all_log_metrics for k, v in d.items()}
 
+        # track all metrics for callbacks
+        self.callback_metrics = {k: v for d in all_callback_metrics for k, v in d.items()}
+
         return 0, grad_norm_dic, all_log_metrics
 
     def __run_evaluation(self, test=False):
@@ -1441,15 +1510,17 @@ class Trainer(TrainerIO):
                                          dataloaders,
                                          max_batches,
                                          test)
-
-            _, progress_bar_metrics, log_metrics = self.__process_output(eval_results)
+            _, prog_bar_metrics, log_metrics, callback_metrics = self.__process_output(eval_results)
 
             # add metrics to prog bar
-            self.__add_tqdm_metrics(progress_bar_metrics)
+            self.__add_tqdm_metrics(prog_bar_metrics)
 
             # log metrics
             self.__log_metrics(log_metrics, {})
 
+            # track metrics for callbacks
+            self.callback_metrics = callback_metrics
+
             # hook
             model.on_post_performance_check()
 
@@ -1460,6 +1531,5 @@ class Trainer(TrainerIO):
 
         # model checkpointing
         if self.proc_rank == 0 and self.checkpoint_callback is not None and not test:
-            print('save callback...')
             self.checkpoint_callback.on_epoch_end(epoch=self.current_epoch,
-                                                  logs=self.__training_tqdm_dict)
+                                                  logs=self.callback_metrics)

pytorch_lightning/trainer/trainer_io.py

@@ -5,7 +5,7 @@ import pdb
 from subprocess import call
 
 import torch
-
+import torch.distributed as dist
 from pytorch_lightning.pt_overrides.override_data_parallel import (
     LightningDistributedDataParallel, LightningDataParallel)
 
@@ -35,6 +35,11 @@ class TrainerIO(object):
         # if script called from hpc resubmit, load weights
         self.restore_hpc_weights_if_needed(model)
 
+        # wait for all models to restore weights
+        if self.use_ddp or self.use_ddp2:
+            # wait for all processes to catch up
+            dist.barrier()
+
     def restore_state_if_checkpoint_exists(self, model):
         # do nothing if there's not dir or callback
         no_ckpt_callback = self.checkpoint_callback is None

setup.py

@@ -14,7 +14,7 @@ from setuptools import setup, find_packages
 # engineer specific practices
 setup(
     name='pytorch-lightning',
-    version='0.5.1',
+    version='0.5.2',
     description='The Keras for ML researchers using PyTorch',
     author='William Falcon',
     author_email='waf2107@columbia.edu',

tests/debug.py

@@ -14,7 +14,9 @@ from torch.utils.data import DataLoader
 from torchvision.datasets import MNIST
 import numpy as np
 import pdb
-from . import test_models
+# from test_models import assert_ok_test_acc, load_model, \
+#     clear_save_dir, get_test_tube_logger, get_hparams, init_save_dir, \
+#     init_checkpoint_callback, reset_seed, set_random_master_port
 
 
 class CoolModel(pl.LightningModule):
@@ -59,56 +61,53 @@ class CoolModel(pl.LightningModule):
     def test_dataloader(self):
         return DataLoader(MNIST('path/to/save', train=False), batch_size=32)
 
-
-def main():
-    """
-    Make sure DDP + AMP continue training correctly
-    :return:
-    """
-    """
-    Make sure DDP2 works
-    :return:
-    """
-    hparams = test_models.get_hparams()
-    model = LightningTestModel(hparams)
-
-    save_dir = test_models.init_save_dir()
-
-    # logger file to get meta
-    logger = test_models.get_test_tube_logger(False)
-    logger.log_hyperparams(hparams)
-    logger.save()
-
-    # logger file to get weights
-    checkpoint = ModelCheckpoint(save_dir)
-
-    trainer_options = dict(
-        show_progress_bar=True,
-        max_nb_epochs=1,
-        train_percent_check=0.4,
-        val_percent_check=0.2,
-        checkpoint_callback=checkpoint,
-        logger=logger,
-        gpus=[0, 1],
-        distributed_backend='dp'
-    )
-
-    # fit model
-    trainer = Trainer(**trainer_options)
-    result = trainer.fit(model)
-
-    # correct result and ok accuracy
-    assert result == 1, 'training failed to complete'
-    pretrained_model = test_models.load_model(logger.experiment, save_dir,
-                                              module_class=LightningTestModel)
-
-    new_trainer = Trainer(**trainer_options)
-    new_trainer.test(pretrained_model)
-
-    # test we have good test accuracy
-    test_models.assert_ok_test_acc(new_trainer)
-    test_models.clear_save_dir()
-
-
-if __name__ == '__main__':
-    main()
+#
+# def main():
+#     reset_seed()
+#     set_random_master_port()
+#
+#     hparams = get_hparams()
+#     model = LightningTestModel(hparams)
+#
+#     save_dir = init_save_dir()
+#
+#     # exp file to get meta
+#     logger = get_test_tube_logger(False)
+#
+#     print(logger.debug)
+#
+#     # exp file to get weights
+#     checkpoint = init_checkpoint_callback(logger)
+#
+#     trainer_options = dict(
+#         show_progress_bar=False,
+#         max_nb_epochs=1,
+#         train_percent_check=0.4,
+#         val_percent_check=0.2,
+#         checkpoint_callback=checkpoint,
+#         logger=logger,
+#         gpus=[0, 1],
+#         distributed_backend='ddp'
+#     )
+#
+#     # fit model
+#     trainer = Trainer(**trainer_options)
+#     result = trainer.fit(model)
+#
+#     exp = logger.experiment
+#     print(os.listdir(exp.get_data_path(exp.name, exp.version)))
+#
+#     # correct result and ok accuracy
+#     assert result == 1, 'training failed to complete'
+#     pretrained_model = load_model(logger.experiment, save_dir,
+#                                   module_class=LightningTestModel)
+#
+#     # run test set
+#     new_trainer = Trainer(**trainer_options)
+#     new_trainer.test(pretrained_model)
+#
+#     # test we have good test accuracy
+#     clear_save_dir()
+#
+# if __name__ == '__main__':
+#     main()

tests/test_logging.py

@@ -7,6 +7,7 @@ import torch
 from pytorch_lightning import Trainer
 
 from pytorch_lightning.testing import LightningTestModel
+from pytorch_lightning.logging import LightningLoggerBase, rank_zero_only
 from .test_models import get_hparams, get_test_tube_logger, init_save_dir, clear_save_dir
 
 RANDOM_SEEDS = list(np.random.randint(0, 10000, 1000))
@@ -134,6 +135,46 @@ def test_mlflow_pickle():
     trainer2.logger.log_metrics({"acc": 1.0})
 
 
+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 reset_seed():
     SEED = RANDOM_SEEDS.pop()
     torch.manual_seed(SEED)

tests/test_models.py

@@ -32,6 +32,7 @@ from pytorch_lightning.logging import TestTubeLogger
 from examples import LightningTemplateModel
 
 # generate a list of random seeds for each test
+RANDOM_PORTS = list(np.random.randint(12000, 19000, 1000))
 ROOT_SEED = 1234
 torch.manual_seed(ROOT_SEED)
 np.random.seed(ROOT_SEED)
@@ -41,6 +42,58 @@ RANDOM_SEEDS = list(np.random.randint(0, 10000, 1000))
 # ------------------------------------------------------------------------
 # TESTS
 # ------------------------------------------------------------------------
+def test_running_test_pretrained_model_ddp():
+    """Verify test() on pretrained model"""
+    if not can_run_gpu_test():
+        return
+
+    reset_seed()
+    set_random_master_port()
+
+    hparams = get_hparams()
+    model = LightningTestModel(hparams)
+
+    save_dir = init_save_dir()
+
+    # exp file to get meta
+    logger = get_test_tube_logger(False)
+
+    # exp file to get weights
+    checkpoint = init_checkpoint_callback(logger)
+
+    trainer_options = dict(
+        show_progress_bar=False,
+        max_nb_epochs=1,
+        train_percent_check=0.4,
+        val_percent_check=0.2,
+        checkpoint_callback=checkpoint,
+        logger=logger,
+        gpus=[0, 1],
+        distributed_backend='ddp'
+    )
+
+    # fit model
+    trainer = Trainer(**trainer_options)
+    result = trainer.fit(model)
+
+    exp = logger.experiment
+    print(os.listdir(exp.get_data_path(exp.name, exp.version)))
+
+    # correct result and ok accuracy
+    assert result == 1, 'training failed to complete'
+    pretrained_model = load_model(logger.experiment, save_dir,
+                                  module_class=LightningTestModel)
+
+    # run test set
+    new_trainer = Trainer(**trainer_options)
+    new_trainer.test(pretrained_model)
+
+    [run_prediction(dataloader, pretrained_model) for dataloader in model.test_dataloader()]
+
+    # test we have good test accuracy
+    clear_save_dir()
+
+
 def test_default_logger_callbacks_cpu_model():
     """
     Test each of the trainer options
@@ -76,10 +129,10 @@ def test_lbfgs_cpu_model():
     trainer_options = dict(
         max_nb_epochs=1,
         gradient_clip_val=1.0,
-        overfit_pct=0.20,
         print_nan_grads=True,
         show_progress_bar=False,
-        train_percent_check=0.2,
+        weights_summary='top',
+        train_percent_check=1.0,
         val_percent_check=0.2
     )
 
@@ -99,9 +152,8 @@ def test_multi_gpu_model_ddp2():
     if not can_run_gpu_test():
         return
 
-    os.environ['MASTER_PORT'] = str(np.random.randint(12000, 19000, 1)[0])
-
     reset_seed()
+    set_random_master_port()
 
     model, hparams = get_model()
     trainer_options = dict(
@@ -110,7 +162,7 @@ def test_multi_gpu_model_ddp2():
         train_percent_check=0.4,
         val_percent_check=0.2,
         gpus=2,
-        print_weights_summary=False,
+        weights_summary=None,
         distributed_backend='ddp2'
     )
 
@@ -141,10 +193,10 @@ def test_dp_resume():
 
     # get logger
     logger = get_test_tube_logger(debug=False)
-    logger.log_hyperparams(hparams)
 
     # exp file to get weights
-    checkpoint = ModelCheckpoint(save_dir)
+    # logger file to get weights
+    checkpoint = init_checkpoint_callback(logger)
 
     # add these to the trainer options
     trainer_options['logger'] = logger
@@ -202,56 +254,6 @@ def test_dp_resume():
     clear_save_dir()
 
 
-def test_running_test_pretrained_model_ddp():
-    """Verify test() on pretrained model"""
-    if not can_run_gpu_test():
-        return
-
-    reset_seed()
-
-    hparams = get_hparams()
-    model = LightningTestModel(hparams)
-
-    save_dir = init_save_dir()
-
-    # exp file to get meta
-    logger = get_test_tube_logger(False)
-    logger.log_hyperparams(hparams)
-    logger.save()
-
-    # exp file to get weights
-    checkpoint = ModelCheckpoint(save_dir)
-
-    trainer_options = dict(
-        show_progress_bar=False,
-        max_nb_epochs=1,
-        train_percent_check=0.4,
-        val_percent_check=0.2,
-        checkpoint_callback=checkpoint,
-        logger=logger,
-        gpus=[0, 1],
-        distributed_backend='ddp'
-    )
-
-    # fit model
-    trainer = Trainer(**trainer_options)
-    result = trainer.fit(model)
-
-    # correct result and ok accuracy
-    assert result == 1, 'training failed to complete'
-    pretrained_model = load_model(logger.experiment, save_dir,
-                                  module_class=LightningTestModel)
-
-    # run test set
-    new_trainer = Trainer(**trainer_options)
-    new_trainer.test(pretrained_model)
-
-    [run_prediction(dataloader, pretrained_model) for dataloader in model.test_dataloader()]
-
-    # test we have good test accuracy
-    clear_save_dir()
-
-
 def test_running_test_after_fitting():
     """Verify test() on fitted model"""
     reset_seed()
@@ -263,11 +265,9 @@ def test_running_test_after_fitting():
 
     # logger file to get meta
     logger = get_test_tube_logger(False)
-    logger.log_hyperparams(hparams)
-    logger.save()
 
     # logger file to get weights
-    checkpoint = ModelCheckpoint(save_dir)
+    checkpoint = init_checkpoint_callback(logger)
 
     trainer_options = dict(
         show_progress_bar=False,
@@ -306,11 +306,9 @@ def test_running_test_without_val():
 
     # logger file to get meta
     logger = get_test_tube_logger(False)
-    logger.log_hyperparams(hparams)
-    logger.save()
 
     # logger file to get weights
-    checkpoint = ModelCheckpoint(save_dir)
+    checkpoint = init_checkpoint_callback(logger)
 
     trainer_options = dict(
         show_progress_bar=False,
@@ -347,11 +345,9 @@ def test_running_test_pretrained_model():
 
     # logger file to get meta
     logger = get_test_tube_logger(False)
-    logger.log_hyperparams(hparams)
-    logger.save()
 
     # logger file to get weights
-    checkpoint = ModelCheckpoint(save_dir)
+    checkpoint = init_checkpoint_callback(logger)
 
     trainer_options = dict(
         show_progress_bar=False,
@@ -394,11 +390,9 @@ def test_running_test_pretrained_model_dp():
 
     # logger file to get meta
     logger = get_test_tube_logger(False)
-    logger.log_hyperparams(hparams)
-    logger.save()
 
     # logger file to get weights
-    checkpoint = ModelCheckpoint(save_dir)
+    checkpoint = init_checkpoint_callback(logger)
 
     trainer_options = dict(
         show_progress_bar=True,
@@ -446,7 +440,7 @@ def test_gradient_accumulation_scheduling():
         assert Trainer(accumulate_grad_batches={1: 2.5, 3: 5})
 
     # test optimizer call freq matches scheduler
-    def optimizer_step(self, epoch_nb, batch_nb, optimizer, optimizer_i):
+    def optimizer_step(self, epoch_nb, batch_nb, optimizer, optimizer_i, second_order_closure=None):
         # only test the first 12 batches in epoch
         if batch_nb < 12:
             if epoch_nb == 0:
@@ -510,8 +504,8 @@ def test_multi_gpu_model_ddp():
     if not can_run_gpu_test():
         return
 
-    os.environ['MASTER_PORT'] = str(np.random.randint(12000, 19000, 1)[0])
     reset_seed()
+    set_random_master_port()
 
     model, hparams = get_model()
     trainer_options = dict(
@@ -646,8 +640,6 @@ def test_no_val_module():
 
     # logger file to get meta
     logger = get_test_tube_logger(False)
-    logger.log_hyperparams(hparams)
-    logger.save()
 
     trainer_options = dict(
         max_nb_epochs=1,
@@ -693,8 +685,6 @@ def test_no_val_end_module():
 
     # logger file to get meta
     logger = get_test_tube_logger(False)
-    logger.log_hyperparams(hparams)
-    logger.save()
 
     trainer_options = dict(
         max_nb_epochs=1,
@@ -829,8 +819,6 @@ def test_cpu_restore_training():
     # logger file to get meta
     test_logger_version = 10
     logger = get_test_tube_logger(False, version=test_logger_version)
-    logger.log_hyperparams(hparams)
-    logger.save()
 
     trainer_options = dict(
         max_nb_epochs=2,
@@ -891,9 +879,8 @@ def test_amp_gpu_ddp():
     if not can_run_gpu_test():
         return
 
-    os.environ['MASTER_PORT'] = str(np.random.randint(12000, 19000, 1)[0])
-
     reset_seed()
+    set_random_master_port()
 
     hparams = get_hparams()
     model = LightningTestModel(hparams)
@@ -923,8 +910,6 @@ def test_cpu_slurm_save_load():
 
     # logger file to get meta
     logger = get_test_tube_logger(False)
-    logger.log_hyperparams(hparams)
-    logger.save()
 
     version = logger.version
 
@@ -961,8 +946,6 @@ def test_cpu_slurm_save_load():
 
     # new logger file to get meta
     logger = get_test_tube_logger(False, version=version)
-    logger.log_hyperparams(hparams)
-    logger.save()
 
     trainer_options = dict(
         max_nb_epochs=1,
@@ -1049,8 +1032,6 @@ def test_model_saving_loading():
 
     # logger file to get meta
     logger = get_test_tube_logger(False)
-    logger.log_hyperparams(hparams)
-    logger.save()
 
     trainer_options = dict(
         max_nb_epochs=1,
@@ -1114,12 +1095,12 @@ def test_amp_gpu_ddp_slurm_managed():
     if not can_run_gpu_test():
         return
 
-    # simulate setting slurm flags
-    os.environ['MASTER_PORT'] = str(np.random.randint(12000, 19000, 1)[0])
-    os.environ['SLURM_LOCALID'] = str(0)
-
     reset_seed()
 
+    # simulate setting slurm flags
+    set_random_master_port()
+    os.environ['SLURM_LOCALID'] = str(0)
+
     hparams = get_hparams()
     model = LightningTestModel(hparams)
 
@@ -1135,11 +1116,9 @@ def test_amp_gpu_ddp_slurm_managed():
 
     # exp file to get meta
     logger = get_test_tube_logger(False)
-    logger.log_hyperparams(hparams)
-    logger.save()
 
     # exp file to get weights
-    checkpoint = ModelCheckpoint(save_dir)
+    checkpoint = init_checkpoint_callback(logger)
 
     # add these to the trainer options
     trainer_options['checkpoint_callback'] = checkpoint
@@ -1350,15 +1329,13 @@ def test_ddp_sampler_error():
     if not can_run_gpu_test():
         return
 
-    os.environ['MASTER_PORT'] = str(np.random.randint(12000, 19000, 1)[0])
-
     reset_seed()
+    set_random_master_port()
 
     hparams = get_hparams()
     model = LightningTestModel(hparams, force_remove_distributed_sampler=True)
 
     logger = get_test_tube_logger(True)
-    logger.save()
 
     trainer = Trainer(
         logger=logger,
@@ -1483,11 +1460,9 @@ def run_gpu_model_test(trainer_options, model, hparams, on_gpu=True):
 
     # logger file to get meta
     logger = get_test_tube_logger(False)
-    logger.log_hyperparams(hparams)
-    logger.save()
 
     # logger file to get weights
-    checkpoint = ModelCheckpoint(save_dir)
+    checkpoint = init_checkpoint_callback(logger)
 
     # add these to the trainer options
     trainer_options['checkpoint_callback'] = checkpoint
@@ -1506,7 +1481,7 @@ def run_gpu_model_test(trainer_options, model, hparams, on_gpu=True):
     # test new model accuracy
     [run_prediction(dataloader, pretrained_model) for dataloader in model.test_dataloader()]
 
-    if trainer.use_ddp:
+    if trainer.use_ddp or trainer.use_ddp2:
         # on hpc this would work fine... but need to hack it for the purpose of the test
         trainer.model = pretrained_model
         trainer.optimizers, trainer.lr_schedulers = pretrained_model.configure_optimizers()
@@ -1557,7 +1532,7 @@ def get_test_tube_logger(debug=True, version=None):
     # set up logger object without actually saving logs
     root_dir = os.path.dirname(os.path.realpath(__file__))
     save_dir = os.path.join(root_dir, 'save_dir')
-    logger = TestTubeLogger(save_dir, name='test_tt_dir', debug=debug, version=version)
+    logger = TestTubeLogger(save_dir, name='lightning_logs', debug=False, version=version)
     return logger
 
 
@@ -1586,10 +1561,11 @@ def load_model(exp, save_dir, module_class=LightningTemplateModel):
 
     # load trained model
     tags_path = exp.get_data_path(exp.name, exp.version)
+    checkpoint_folder = os.path.join(tags_path, 'checkpoints')
     tags_path = os.path.join(tags_path, 'meta_tags.csv')
 
-    checkpoints = [x for x in os.listdir(save_dir) if '.ckpt' in x]
-    weights_dir = os.path.join(save_dir, checkpoints[0])
+    checkpoints = [x for x in os.listdir(checkpoint_folder) if '.ckpt' in x]
+    weights_dir = os.path.join(checkpoint_folder, checkpoints[0])
 
     trained_model = module_class.load_from_metrics(weights_path=weights_dir,
                                                    tags_csv=tags_path)
@@ -1654,5 +1630,18 @@ def reset_seed():
     np.random.seed(SEED)
 
 
+def set_random_master_port():
+    port = RANDOM_PORTS.pop()
+    os.environ['MASTER_PORT'] = str(port)
+
+
+def init_checkpoint_callback(logger):
+    exp = logger.experiment
+    exp_path = exp.get_data_path(exp.name, exp.version)
+    ckpt_dir = os.path.join(exp_path, 'checkpoints')
+    checkpoint = ModelCheckpoint(ckpt_dir)
+    return checkpoint
+
+
 if __name__ == '__main__':
     pytest.main([__file__])