Merge branch 'master' of https://github.com/williamFalcon/pytorch-lightning

.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. 
 

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.   
@@ -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)

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 
 

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)
+```

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)

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,

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)

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

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):
@@ -177,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
@@ -187,8 +199,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 +245,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 +305,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

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")

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

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:

pytorch_lightning/root_module/memory.py

@@ -3,11 +3,13 @@ Generates a summary of a model's layers and dimensionality
 '''
 
 import gc
+import os
 import subprocess
 
 import numpy as np
 import pandas as pd
 import torch
+import logging
 
 
 class ModelSummary(object):
@@ -166,7 +168,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
 
@@ -198,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
 
@@ -224,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
 
 

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
@@ -10,6 +12,8 @@ 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
+from pytorch_lightning.pt_overrides.override_data_parallel import LightningDistributedDataParallel
 
 
 class LightningModule(GradInformation, ModelIO, ModelHooks):
@@ -47,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
@@ -89,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)
@@ -240,12 +319,16 @@ 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():
             param.requires_grad = False
 
+        self.eval()
+
     def unfreeze(self):
         for param in self.parameters():
             param.requires_grad = True
+
+        self.train()

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 = """

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

pytorch_lightning/trainer/ddp_mixin.py

@@ -1,11 +1,10 @@
 import os
 import re
 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:
@@ -59,7 +58,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 +106,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):
         """
@@ -144,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
@@ -166,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
@@ -177,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]

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)

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])

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

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:

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)

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

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'

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():

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():