release v0.4.9

* load from metrics defaults to CPU

* load from metrics defaults to CPU

* load from metrics defaults to CPU

.github/PULL_REQUEST_TEMPLATE.md

@@ -0,0 +1,16 @@
+# Before submitting
+
+- Was this discussed/approved via a Github issue? (no need for typos, doc improvements)
+- Did you read the [contributor guideline](https://github.com/williamFalcon/pytorch-lightning/blob/master/.github/CONTRIBUTING.md)?
+- Did you make sure to update the docs?   
+- Did you write any new necessary tests?  
+
+## What does this PR do?
+Fixes # (issue).
+
+## PR review    
+Anyone in the community is free to review the PR once the tests have passed.     
+If we didn't discuss your PR in Github issues there's a high chance it will not be merged.   
+
+## Did you have fun?
+Make sure you had fun coding 🙃

.travis.yml

@@ -8,8 +8,6 @@
 # this file is *not* meant to cover or endorse the use of travis, but rather to
 # help confirm pull requests to this project.
 
-dist: xenial  # Ubuntu 16.04
-
 env:
   global:
     - DISPLAY=""
@@ -18,10 +16,36 @@ language: python
 
 matrix:
   include:
-    - python: 3.6
+    - os: linux
+      dist: xenial  # Ubuntu 16.04
+      python: 3.6
       env: TOXENV=py36
-    - python: 3.7
+    - os: linux
+      dist: bionic  # Ubuntu 18.04
+      python: 3.6
+      env: TOXENV=py36
+    - os: linux
+      dist: bionic  # Ubuntu 18.04
+      python: 3.7
       env: TOXENV=py37
+    - os: osx
+      osx_image: xcode9.4
+      language: generic
+      env: TOXENV=py36
+      addons:
+        homebrew:
+          # update: true
+          packages: python3
+      before_install:
+        - pip3 install virtualenv
+        - virtualenv -p python3 ~/venv
+        - source ~/venv/bin/activate
+#    - os: windows
+#      language: minimal
+#      before_install:
+#        - choco install python3
+#        - export PATH="/c/Python37:/c/Python37/Scripts:$PATH"
+#      env: TOXENV=py37
 
 # See http://docs.travis-ci.com/user/caching/#pip-cache
 cache: pip

MANIFEST.in

@@ -13,11 +13,16 @@ include LICENSE
 exclude *.sh
 exclude *.toml
 exclude *.svg    
-recursive-include examples *.py
 recursive-include pytorch_lightning *.py
 
+# include examples
+recursive-include examples *.py
+recursive-include examples *.md
+recursive-include examples *.sh
+
 # exclude tests from package
 recursive-exclude tests *
+recursive-exclude site *
 exclude tests
 
 # Exclude the documentation files

README.md

@@ -4,7 +4,7 @@
 
 # PyTorch Lightning
 
-**The PyTorch Keras for ML researchers. More control. Less boilerplate.**
+**The lightweight PyTorch wrapper for ML researchers. Scale your models. Write less boilerplate.**
 
 
 [![PyPI Status](https://badge.fury.io/py/pytorch-lightning.svg)](https://badge.fury.io/py/pytorch-lightning)
@@ -43,7 +43,8 @@ Every research project starts the same, a model, a training loop, validation loo
 
 Lightning sets up all the boilerplate state-of-the-art training for you so you can focus on the research.   
 
----   
+---
+ 
 ## README Table of Contents        
 - [How do I use it](https://github.com/williamFalcon/pytorch-lightning#how-do-i-do-use-it)     
 - [What lightning automates](https://github.com/williamFalcon/pytorch-lightning#what-does-lightning-control-for-me)    
@@ -57,7 +58,8 @@ Lightning sets up all the boilerplate state-of-the-art training for you so you c
 - [Asking for help](https://github.com/williamFalcon/pytorch-lightning#asking-for-help)
 - [FAQ](https://github.com/williamFalcon/pytorch-lightning#faq)    
 
----   
+---
+
 ## How do I do use it?   
 Think about Lightning as refactoring your research code instead of using a new framework. The research code goes into a [LightningModule]((https://williamfalcon.github.io/pytorch-lightning/LightningModule/RequiredTrainerInterface/)) which you fit using a Trainer.    
 
@@ -122,7 +124,6 @@ class CoolSystem(pl.LightningModule):
         # OPTIONAL
         return DataLoader(MNIST(os.getcwd(), train=True, download=True, transform=transforms.ToTensor()), batch_size=32)
 ```
-
 2. Fit with a [trainer](https://williamfalcon.github.io/pytorch-lightning/Trainer/)    
 ```python
 from pytorch_lightning import Trainer
@@ -134,7 +135,7 @@ trainer = Trainer()
 trainer.fit(model)   
 ```     
 
-Or with tensorboard logger and some options turned on such as multi-gpu, etc...    
+Or with tensorboard logger and some options turned on such as multi-gpu, etc...
 ```python   
 from test_tube import Experiment    
 
@@ -159,6 +160,11 @@ print('View tensorboard logs by running\ntensorboard --logdir %s' % os.getcwd())
 print('and going to http://localhost:6006 on your browser')
 ```    
 
+When you're all done you can even run the test set separately.   
+```python
+trainer.test()
+```
+
 ## What does lightning control for me?   
 
 Everything in gray!    
@@ -252,10 +258,10 @@ 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    
+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
+from pytorch_lightning import Trainer
 
 exp = Experiment(save_dir='/some/path')
 trainer = Trainer(experiment=exp)
@@ -270,18 +276,18 @@ tensorboard --logdir /some/path
 ## Lightning automates all of the following ([each is also configurable](https://williamfalcon.github.io/pytorch-lightning/Trainer/)):
 
 
-###### Checkpointing    
+#### Checkpointing    
 
 - [Model saving](https://williamfalcon.github.io/pytorch-lightning/Trainer/Checkpointing/#model-saving)
 - [Model loading](https://williamfalcon.github.io/pytorch-lightning/LightningModule/methods/#load-from-metrics) 
 - [Restoring training session](https://williamfalcon.github.io/pytorch-lightning/Trainer/Checkpointing/#restoring-training-session)
 
-###### Computing cluster (SLURM)    
+#### Computing cluster (SLURM)    
 
 - [Running grid search on a cluster](https://williamfalcon.github.io/pytorch-lightning/Trainer/SLURM%20Managed%20Cluster#running-grid-search-on-a-cluster) 
 - [Walltime auto-resubmit](https://williamfalcon.github.io/pytorch-lightning/Trainer/SLURM%20Managed%20Cluster#walltime-auto-resubmit)   
 
-###### Debugging  
+#### Debugging  
 
 - [Fast dev run](https://williamfalcon.github.io/pytorch-lightning/Trainer/debugging/#fast-dev-run)
 - [Inspect gradient norms](https://williamfalcon.github.io/pytorch-lightning/Trainer/debugging/#inspect-gradient-norms)
@@ -292,7 +298,7 @@ tensorboard --logdir /some/path
 - [Print input and output size of every module in system](https://williamfalcon.github.io/pytorch-lightning/LightningModule/properties/#example_input_array)
 
 
-###### Distributed training    
+#### Distributed training    
 
 - [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)
@@ -301,7 +307,7 @@ tensorboard --logdir /some/path
 - [Self-balancing architecture](https://williamfalcon.github.io/pytorch-lightning/Trainer/Distributed%20training/#self-balancing-architecture)
 
 
-###### Experiment Logging   
+#### Experiment Logging   
 
 - [Display metrics in progress bar](https://williamfalcon.github.io/pytorch-lightning/Trainer/Logging/#display-metrics-in-progress-bar)
 - [Log metric row every k batches](https://williamfalcon.github.io/pytorch-lightning/Trainer/Logging/#log-metric-row-every-k-batches)
@@ -311,7 +317,7 @@ tensorboard --logdir /some/path
 - [Snapshot code for a training run](https://williamfalcon.github.io/pytorch-lightning/Trainer/Logging/#snapshot-code-for-a-training-run) 
 - [Write logs file to csv every k batches](https://williamfalcon.github.io/pytorch-lightning/Trainer/Logging/#write-logs-file-to-csv-every-k-batches)
 
-###### Training loop    
+#### Training loop    
 
 - [Accumulate gradients](https://williamfalcon.github.io/pytorch-lightning/Trainer/Training%20Loop/#accumulated-gradients)
 - [Force training for min or max epochs](https://williamfalcon.github.io/pytorch-lightning/Trainer/Training%20Loop/#force-training-for-min-or-max-epochs)
@@ -323,7 +329,7 @@ tensorboard --logdir /some/path
 - [Set how much of the training set to check (1-100%)](https://williamfalcon.github.io/pytorch-lightning/Trainer/Training%20Loop/#set-how-much-of-the-training-set-to-check)
 - [Step optimizers at arbitrary intervals](https://williamfalcon.github.io/pytorch-lightning/Trainer/hooks/#optimizer_step)
 
-###### Validation loop    
+#### Validation loop    
 
 - [Check validation every n epochs](https://williamfalcon.github.io/pytorch-lightning/Trainer/Validation%20loop/#check-validation-every-n-epochs)
 - [Hooks](https://williamfalcon.github.io/pytorch-lightning/Trainer/hooks/)
@@ -332,11 +338,13 @@ tensorboard --logdir /some/path
 - [Set validation check frequency within 1 training epoch](https://williamfalcon.github.io/pytorch-lightning/Trainer/Validation%20loop/#set-validation-check-frequency-within-1-training-epoch)
 - [Set the number of validation sanity steps](https://williamfalcon.github.io/pytorch-lightning/Trainer/Validation%20loop/#set-the-number-of-validation-sanity-steps)
 
+#### Testing loop  
+- [Run test set](https://williamfalcon.github.io/pytorch-lightning/Trainer/Testing%20loop/)  
 
 ## Demo
 ```bash
 # install lightning
-pip install pytorch-lightning
+pip install pytorch_lightning
 
 # clone lightning for the demo
 git clone https://github.com/williamFalcon/pytorch-lightning.git
@@ -371,7 +379,7 @@ If you have any questions, feel free to:
 
 If no one replies to you quickly enough, feel free to post the stackoverflow link to our Gitter chat!   
 
-To chat with the rest of us visit our [gitter channel](https://gitter.im/PyTorch-Lightning/community?utm_source=share-link&utm_medium=link&utm_campaign=share-link)!     
+To chat with the rest of us visit our [gitter channel](https://gitter.im/PyTorch-Lightning/community)!     
 
 ---   
 ## FAQ    
@@ -397,22 +405,36 @@ Nope.
 Nope. Please use anaconda or miniconda.    
 
 **Which PyTorch versions do you support?**    
-##### PyTorch 1.1.0       
-```bash    
-# install pytorch 1.1.0 using the official instructions   
+- **PyTorch 1.1.0**       
+    ```bash    
+    # install pytorch 1.1.0 using the official instructions   
+    
+    # install test-tube 0.6.7.6 which supports 1.1.0   
+    pip install test-tube==0.6.7.6   
+    
+    # install latest Lightning version without upgrading deps    
+    pip install -U --no-deps pytorch-lightning
+    ```     
+- **PyTorch 1.2.0**
+    Install via pip as normal   
 
-# install test-tube 0.6.7.6 which supports 1.1.0   
-pip install test-tube==0.6.7.6   
+## Custom installation
 
-# install latest Lightning version without upgrading deps    
-pip install -U --no-deps pytorch-lightning
-```     
+### Bleeding edge
 
-##### PyTorch 1.2.0   
-Install via pip as normal   
+If you can't wait for the next release, install the most up to date code with:
+* using GIT (locally clone whole repo with full history)
+    ```bash
+    pip install git+https://github.com/williamFalcon/pytorch-lightning.git@master --upgrade
+    ```
+* using instant zip (last state of the repo without git history)
+    ```bash
+    pip install https://github.com/williamFalcon/pytorch-lightning/archive/master.zip --upgrade
+    ```
 
-## Bleeding edge
-If you can't wait for the next release, install the most up to date code with:  
+### Any release installation
+
+You can also install any past release from this repository:
 ```bash
-pip install git+https://github.com/williamFalcon/pytorch-lightning.git@master --upgrade
+pip install https://github.com/williamFalcon/pytorch-lightning/archive/0.4.4.zip --upgrade
 ```

docs/LightningModule/RequiredTrainerInterface.md

@@ -16,7 +16,9 @@ Otherwise, to Define a Lightning Module, implement the following methods:
 **Optional**:   
 
 - [validation_step](RequiredTrainerInterface.md#validation_step)    
-- [validation_end](RequiredTrainerInterface.md#validation_end)    
+- [validation_end](RequiredTrainerInterface.md#validation_end) 
+- [test_step](RequiredTrainerInterface.md#test_step)    
+- [test_end](RequiredTrainerInterface.md#test_end) 
 - [val_dataloader](RequiredTrainerInterface.md#val_dataloader)    
 - [test_dataloader](RequiredTrainerInterface.md#test_dataloader)    
 - [on_save_checkpoint](RequiredTrainerInterface.md#on_save_checkpoint)    
@@ -63,6 +65,17 @@ class CoolModel(pl.LightningModule):
         avg_loss = torch.stack([x['val_loss'] for x in outputs]).mean()
         return {'avg_val_loss': avg_loss}
 
+    def test_step(self, batch, batch_nb):
+        # OPTIONAL
+        x, y = batch
+        y_hat = self.forward(x)
+        return {'test_loss': F.cross_entropy(y_hat, y)}
+
+    def test_end(self, outputs):
+        # OPTIONAL
+        avg_loss = torch.stack([x['test_loss'] for x in outputs]).mean()
+        return {'avg_test_loss': avg_loss}
+
     def configure_optimizers(self):
         # REQUIRED
         return [torch.optim.Adam(self.parameters(), lr=0.02)]
@@ -80,6 +93,7 @@ class CoolModel(pl.LightningModule):
     @pl.data_loader
     def test_dataloader(self):
         # OPTIONAL
+        # can also return a list of test dataloaders
         return DataLoader(MNIST(os.getcwd(), train=False, download=True, transform=transforms.ToTensor()), batch_size=32)
 ```
 ---   
@@ -225,17 +239,18 @@ the [optimizer_step](https://williamfalcon.github.io/pytorch-lightning/Trainer/h
 ### validation_step
 
 ``` {.python}
+# if you have one val dataloader:
 def validation_step(self, data_batch, batch_nb)   
 
-# if have multiple val dataloaders:  
+# if you have multiple val dataloaders:  
 def validation_step(self, data_batch, batch_nb, dataloader_idx)
 ```
 **OPTIONAL**    
-If you don't need to validate you don't need to implement this method.    
+If you don't need to validate you don't need to implement this method. In this step you'd normally generate examples or calculate anything of interest such as accuracy. 
 
-In this step you'd normally generate examples or calculate anything of interest such as accuracy. 
+When the validation_step is called, the model has been put in eval mode and PyTorch gradients have been disabled. At the end of validation, model goes back to training mode and gradients are enabled.
 
-The dict you return here will be available in the validation_end method. 
+The dict you return here will be available in the `validation_end` method. 
 
 **Params**   
 
@@ -256,11 +271,11 @@ The dict you return here will be available in the validation_end method.
 ``` {.python}
 # CASE 1: A single validation dataset
 def validation_step(self, data_batch, batch_nb):
-    x, y, z = data_batch
+    x, y = data_batch
     
     # implement your own
     out = self.forward(x)
-    loss = self.loss(out, x)
+    loss = self.loss(out, y)
     
     # log 6 example images
     # or generated text... or whatever
@@ -301,7 +316,7 @@ def validation_end(self, outputs)
 ```   
 If you didn't define a validation_step, this won't be called.       
 
-Called at the end of the validation loop with the output of each validation_step.  Called once per validation dataset.   
+Called at the end of the validation loop with the outputs of validation_step.
 
 The outputs here are strictly for the progress bar. If you don't need to display anything, don't return anything.    
 
@@ -309,7 +324,7 @@ The outputs here are strictly for the progress bar. If you don't need to display
 
 | Param  | description  |
 |---|---|
-|  outputs | List of outputs you defined in validation_step |
+|  outputs | List of outputs you defined in validation_step, or if there are multiple dataloaders, a list containing a list of outputs for each dataloader |
 
 **Return**   
 
@@ -319,6 +334,8 @@ The outputs here are strictly for the progress bar. If you don't need to display
 
 **Example**
 
+With a single dataloader
+
 ``` {.python}
 def validation_end(self, outputs):
     """
@@ -338,6 +355,171 @@ def validation_end(self, outputs):
     return tqdm_dic
 ```
 
+With multiple dataloaders, `outputs` will be a list of lists. The outer list contains
+one entry per dataloader, while the inner list contains the individual outputs of 
+each validation step for that dataloader.
+
+``` {.python}
+def validation_end(self, outputs):
+    """
+    Called at the end of validation to aggregate outputs
+    :param outputs: list of list of individual outputs of each validation step
+    :return:
+    """
+    val_loss_mean = 0
+    val_acc_mean = 0
+    i = 0
+    for dataloader_outputs in outputs:
+        for output in dataloader_outputs:
+            val_loss_mean += output['val_loss']
+            val_acc_mean += output['val_acc']
+            i += 1
+
+    val_loss_mean /= i
+    val_acc_mean /= i
+    tqdm_dic = {'val_loss': val_loss_mean.item(), 'val_acc': val_acc_mean.item()}
+    return tqdm_dic
+```
+
+### test_step
+
+``` {.python}
+# if you have one test dataloader:
+def test_step(self, data_batch, batch_nb)   
+
+# if you have multiple test dataloaders:  
+def test_step(self, data_batch, batch_nb, dataloader_idx)
+```
+**OPTIONAL**    
+If you don't need to test you don't need to implement this method. In this step you'd normally generate examples or calculate anything of interest such as accuracy.   
+
+When the validation_step is called, the model has been put in eval mode and PyTorch gradients have been disabled. At the end of validation, model goes back to training mode and gradients are enabled.
+
+The dict you return here will be available in the `test_end` method. 
+
+This function is used when you execute `trainer.test()`.
+
+**Params**   
+
+| Param  | description  |
+|---|---|
+|  data_batch | The output of your dataloader. A tensor, tuple or list  |
+|  batch_nb | Integer displaying which batch this is  |
+|  dataloader_i | Integer displaying which dataloader this is (only if multiple test datasets used)  |
+
+**Return**   
+
+| Return  | description  | optional |
+|---|---|---|   
+|  dict | Dict or OrderedDict with metrics to display in progress bar. All keys must be tensors. | Y |
+
+**Example**
+
+``` {.python}
+# CASE 1: A single test dataset
+def test_step(self, data_batch, batch_nb):
+    x, y = data_batch
+    
+    # implement your own
+    out = self.forward(x)
+    loss = self.loss(out, y)
+    
+    # calculate acc
+    labels_hat = torch.argmax(out, dim=1)
+    test_acc = torch.sum(y == labels_hat).item() / (len(y) * 1.0)
+    
+    # all optional...
+    # return whatever you need for the collation function test_end
+    output = OrderedDict({
+        'test_loss': loss_test,
+        'test_acc': torch.tensor(test_acc), # everything must be a tensor
+    })
+    
+    # return an optional dict
+    return output
+```   
+
+If you pass in multiple test datasets, test_step will have an additional argument.
+
+```python
+# CASE 2: multiple test datasets
+def test_step(self, data_batch, batch_nb, dataset_idx):
+    # dataset_idx tells you which dataset this is.   
+```   
+
+The ```dataset_idx``` corresponds to the order of datasets returned in ```test_dataloader```.    
+
+--- 
+### test_end
+
+``` {.python}
+def test_end(self, outputs)
+```   
+If you didn't define a test_step, this won't be called.       
+
+Called at the end of the test step with the output of each test_step.
+
+The outputs here are strictly for the progress bar. If you don't need to display anything, don't return anything.    
+
+**Params**    
+
+| Param  | description  |
+|---|---|
+|  outputs | List of outputs you defined in test_step, or if there are multiple dataloaders, a list containing a list of outputs for each dataloader |
+
+**Return**   
+
+| Return  | description  | optional |
+|---|---|---|   
+|  dict | Dict of OrderedDict with metrics to display in progress bar | Y |
+
+**Example**
+
+``` {.python}
+def test_end(self, outputs):
+    """
+    Called at the end of test to aggregate outputs
+    :param outputs: list of individual outputs of each test step
+    :return:
+    """
+    test_loss_mean = 0
+    test_acc_mean = 0
+    for output in outputs:
+        test_loss_mean += output['test_loss']
+        test_acc_mean += output['test_acc']
+
+    test_loss_mean /= len(outputs)
+    test_acc_mean /= len(outputs)
+    tqdm_dic = {'test_loss': test_loss_mean.item(), 'test_acc': test_acc_mean.item()}
+    return tqdm_dic
+```
+
+With multiple dataloaders, `outputs` will be a list of lists. The outer list contains
+one entry per dataloader, while the inner list contains the individual outputs of 
+each validation step for that dataloader.
+
+``` {.python}
+def test_end(self, outputs):
+    """
+    Called at the end of test to aggregate outputs
+    :param outputs: list of individual outputs of each test step
+    :return:
+    """
+    test_loss_mean = 0
+    test_acc_mean = 0
+    i = 0
+    for dataloader_outputs in outputs:
+        for output in dataloader_outputs:
+            test_loss_mean += output['test_loss']
+            test_acc_mean += output['test_acc']
+            i += 1
+
+    test_loss_mean /= i 
+    test_acc_mean /= i
+    tqdm_dic = {'test_loss': test_loss_mean.item(), 'test_acc': test_acc_mean.item()}
+    return tqdm_dic
+```
+
 --- 
 ### on_save_checkpoint 
 

docs/Trainer/Distributed training.md

@@ -10,10 +10,13 @@ For multi-node training you must use DistributedDataParallel.
 
 You can toggle between each mode by setting this flag.
 ``` {.python}
-# DEFAULT uses DataParallel
+# DEFAULT (when using single GPU or no GPUs)
+trainer = Trainer(distributed_backend=None)
+
+# Change to DataParallel (gpus > 1)
 trainer = Trainer(distributed_backend='dp')
 
-# change to distributed data parallel
+# change to distributed data parallel (gpus > 1)
 trainer = Trainer(distributed_backend='ddp')
 ```
 
@@ -32,12 +35,24 @@ Below are the possible configurations we support.
 
 | 1 GPU  | 1+ GPUs  | DP  | DDP  | 16-bit  | command |
 |---|---|---|---|---|---|
-| Y  |   |   |   |  | ```Trainer(gpus=[0])``` |
-| Y  |   |   |   | Y | ```Trainer(gpus=[0], use_amp=True)``` |
-|   | Y | Y |   |   | ```Trainer(gpus=[0, ...])``` |
-|   | Y |  | Y  |  | ```Trainer(gpus=[0, ...], distributed_backend='ddp')``` |
-|   | Y |  | Y  | Y | ```Trainer(gpus=[0, ...], distributed_backend='ddp', use_amp=True)``` |
+| Y  |   |   |   |  | ```Trainer(gpus=1)``` |
+| Y  |   |   |   | Y | ```Trainer(gpus=1, use_amp=True)``` |
+|   | Y | Y |   |   | ```Trainer(gpus=k)``` |
+|   | Y |  | Y  |  | ```Trainer(gpus=k, distributed_backend='ddp')``` |
+|   | Y |  | Y  | Y | ```Trainer(gpus=k, distributed_backend='ddp', use_amp=True)``` |
 
+You also have the option of specifying which GPUs to use by passing a list:   
+
+```python
+# DEFAULT (int)
+Trainer(gpus=k)  
+
+# You specify which GPUs (don't use if running on cluster)  
+Trainer(gpus=[0, 1])  
+
+# can also be a string
+Trainer(gpus='0, 1')
+```
 
 ---
 #### CUDA flags   
@@ -49,6 +64,9 @@ Lightning sets these for you automatically, there's NO NEED to do this yourself.
 # os.environ["CUDA_VISIBLE_DEVICES"] = "0"
 ```
 
+However, when using a cluster, Lightning will NOT set these flags (and you should not either). 
+SLURM will set these for you.   
+
 ---
 #### 16-bit mixed precision
 16 bit precision can cut your memory footprint by half. If using volta architecture GPUs it can give a dramatic training speed-up as well.    
@@ -70,7 +88,7 @@ trainer = Trainer(amp_level='O2', use_amp=False)
 Make sure you're on a GPU machine. 
 ```python
 # DEFAULT
-trainer = Trainer(gpus=[0])
+trainer = Trainer(gpus=1)
 ```
 
 ---
@@ -78,11 +96,11 @@ trainer = Trainer(gpus=[0])
 Make sure you're on a GPU machine. You can set as many GPUs as you want.
 In this setting, the model will run on all 8 GPUs at once using DataParallel under the hood.
 ```python
-# to use DataParallel (default)
-trainer = Trainer(gpus=[0,1,2,3,4,5,6,7], distributed_backend='dp')
+# to use DataParallel
+trainer = Trainer(gpus=8, distributed_backend='dp')
 
 # RECOMMENDED use DistributedDataParallel
-trainer = Trainer(gpus=[0,1,2,3,4,5,6,7], distributed_backend='ddp')
+trainer = Trainer(gpus=8, distributed_backend='ddp')
 ```
 
 ---
@@ -90,7 +108,7 @@ trainer = Trainer(gpus=[0,1,2,3,4,5,6,7], distributed_backend='ddp')
 Multi-node training is easily done by specifying these flags.
 ```python
 # train on 12*8 GPUs
-trainer = Trainer(gpus=[0,1,2,3,4,5,6,7], nb_gpu_nodes=12)
+trainer = Trainer(gpus=8, nb_gpu_nodes=12)
 ```
 
 In addition, make sure to set up your SLURM job correctly via the [SlurmClusterObject](https://williamfalcon.github.io/test-tube/hpc/SlurmCluster/). In particular, specify the number of tasks per node correctly.

docs/Trainer/Logging.md

@@ -14,6 +14,14 @@ Every k batches lightning will make an entry in the metrics log
 ``` {.python}
 # DEFAULT (ie: save a .csv log file every 10 batches)
 trainer = Trainer(add_log_row_interval=10)
+```   
+
+---
+#### Log metric row every k batches 
+Logs GPU memory when metrics are logged.   
+``` {.python}
+# DEFAULT
+trainer = Trainer(log_gpu_memory=False)
 ```
 
 ---

docs/Trainer/SLURM Managed Cluster.md

@@ -1,8 +1,10 @@
 Lightning supports model training on a cluster managed by SLURM in the following cases:    
 
-1. Training on single or multi-cpus only.
-2. Training on single or multi-gpus on the same node.
-3. Coming SOON: Training across multiple nodes.
+1. Training on a single cpu or single GPU.
+2. Train on multiple GPUs on the same node using DataParallel or DistributedDataParallel
+3. Training across multiple GPUs on multiple different nodes via DistributedDataParallel.
+
+**Note: A node means a machine with multiple GPUs**
 
 ---
 #### Running grid search on a cluster
@@ -55,8 +57,8 @@ cluster.memory_mb_per_node = 10000
 cluster.job_time = '10:00'
 ```
 
-(3). Give trainer the cluster_manager in your main function:    
-
+(3). Make a main function with your model and trainer. Each job will call this function with a particular
+hparams configuration.    
 ```{.python}
 from pytorch_lightning import Trainer
 
@@ -66,12 +68,12 @@ def train_fx(trial_hparams, cluster_manager, _):
     my_model = MyLightningModel()
     
     # give the trainer the cluster object
-    trainer = Trainer(cluster=cluster_manager)
+    trainer = Trainer()
     trainer.fit(my_model)
 
 ```
 
-(4). Start the grid search     
+(3). Start the grid/random search     
 ```{.python}
 # run the models on the cluster
 cluster.optimize_parallel_cluster_gpu(
@@ -81,24 +83,22 @@ cluster.optimize_parallel_cluster_gpu(
     job_display_name='my_exp')
 ```
 
-That's it! The SlurmCluster object will automatically checkpoint the lightning model and resubmit if it runs into the walltime!
-
-
 ---
 #### Walltime auto-resubmit
-Lightning automatically resubmits jobs when they reach the walltime. You get this behavior for free if you give lightning
-a slurm cluster object.
+Lightning automatically resubmits jobs when they reach the walltime. Make sure to set the SIGUSR1 signal in 
+your SLURM script.   
 
-```{.python}
-def my_main_fx(hparams, slurm_manager, _):
-    trainer = Trainer(cluster=slurm_manager)
+```bash
+# 90 seconds before training ends
+#SBATCH --signal=SIGUSR1@90
 ``` 
 
-(See the grid search example above for cluster configuration).
-With this feature lightning will:    
+When lightning receives the SIGUSR1 signal it will:
+1. save a checkpoint with 'hpc_ckpt' in the name.
+2. resubmit the job using the SLURM_JOB_ID  
+
+When the script starts again, Lightning will:
+1. search for a 'hpc_ckpt' checkpoint. 
+2. restore the model, optimizers, schedulers, epoch, etc...   
 
-1. automatically checkpoint the model
-2. checkpoint the trainer session
-3. resubmit a continuation job.
-4. load the checkpoint and trainer session in the new model
 

docs/Trainer/Testing loop.md

@@ -0,0 +1,31 @@
+To ensure you don't accidentally use test data to guide training decisions Lightning makes running the test set deliberate.   
+
+---
+#### test
+You have two options to run the test set. 
+First case is where you test right after a full training routine.
+``` {.python}
+# run full training
+trainer.fit(model)
+
+# run test set
+trainer.test()
+```   
+
+Second case is where you load a model and run the test set   
+```{.python}
+model = MyLightningModule.load_from_metrics(
+    weights_path='/path/to/pytorch_checkpoint.ckpt',
+    tags_csv='/path/to/test_tube/experiment/version/meta_tags.csv',
+    on_gpu=True,
+    map_location=None
+)
+
+# init trainer with whatever options
+trainer = Trainer(...)
+    
+# test (pass in the model)
+trainer.test(model)
+```
+In this second case, the options you pass to trainer will be used when running the test set (ie: 16-bit, dp, ddp, etc...)  
+ 

docs/Trainer/index.md

@@ -78,3 +78,8 @@ But of course the fun is in all the advanced things it can do:
 - [Set how much of the test set to check](https://williamfalcon.github.io/pytorch-lightning/Trainer/Validation%20loop/#set-how-much-of-the-test-set-to-check)
 - [Set validation check frequency within 1 training epoch](https://williamfalcon.github.io/pytorch-lightning/Trainer/Validation%20loop/#set-validation-check-frequency-within-1-training-epoch)
 - [Set the number of validation sanity steps](https://williamfalcon.github.io/pytorch-lightning/Trainer/Validation%20loop/#set-the-number-of-validation-sanity-steps)
+
+
+**Testing loop**    
+
+- [Run test set](https://williamfalcon.github.io/pytorch-lightning/Trainer/Testing%20loop/)  

docs/index.md

@@ -137,3 +137,5 @@ Notice a few things about this flow:
 - [Set validation check frequency within 1 training epoch](https://williamfalcon.github.io/pytorch-lightning/Trainer/Validation%20loop/#set-validation-check-frequency-within-1-training-epoch)
 - [Set the number of validation sanity steps](https://williamfalcon.github.io/pytorch-lightning/Trainer/Validation%20loop/#set-the-number-of-validation-sanity-steps)
 
+###### Testing loop  
+- [Run test set](https://williamfalcon.github.io/pytorch-lightning/Trainer/Testing%20loop/)  

examples/new_project_templates/lightning_module_template.py

@@ -240,15 +240,15 @@ class LightningTemplateModel(LightningModule):
         parser.add_argument('--out_features', default=10, type=int)
         # use 500 for CPU, 50000 for GPU to see speed difference
         parser.add_argument('--hidden_dim', default=50000, type=int)
-        parser.opt_list('--drop_prob', default=0.2, options=[0.2, 0.5], type=float, tunable=False)
+        parser.opt_list('--drop_prob', default=0.2, options=[0.2, 0.5], type=float, tunable=True)
+        parser.opt_list('--learning_rate', default=0.001 * 8, type=float,
+                        options=[0.0001, 0.0005, 0.001],
+                        tunable=True)
 
         # data
         parser.add_argument('--data_root', default=os.path.join(root_dir, 'mnist'), type=str)
 
         # training params (opt)
-        parser.opt_list('--learning_rate', default=0.001 * 8, type=float,
-                        options=[0.0001, 0.0005, 0.001, 0.005],
-                        tunable=False)
         parser.opt_list('--optimizer_name', default='adam', type=str,
                         options=['adam'], tunable=False)
 

examples/new_project_templates/multi_node_examples/README.md

@@ -0,0 +1,107 @@
+# Multi-node examples
+Use these templates for multi-node training. 
+The main complexity around cluster training is how you submit the SLURM jobs.  
+
+## Test-tube   
+Lightning uses test-tube to submit SLURM jobs and to run hyperparameter searches on a cluster.  
+
+To run a hyperparameter search, we normally add the values to search to the Hyperparameter optimizer 
+```python
+from test_tube import HyperOptArgumentParser
+
+parser = HyperOptArgumentParser(strategy='grid_search')
+parser.opt_list('--drop_prob', default=0.2, options=[0.2, 0.5], type=float, tunable=True)
+parser.opt_list('--learning_rate', default=0.001, type=float,
+                        options=[0.0001, 0.0005, 0.001],
+                        tunable=True)
+                        
+# give your model a chance to add its own parameters
+parser = LightningTemplateModel.add_model_specific_args(parent_parser, root_dir)
+
+# parse args
+hyperparams = parser.parse_args()
+```
+
+The above sets up a grid search on learning rate and drop probability. You can now add this object to the 
+cluster object to perform the grid search:   
+```python
+cluster = SlurmCluster(
+    hyperparam_optimizer=hyperparams,
+    log_path='/path/to/log/slurm/files',
+)
+
+# ... configure cluster options
+
+# run grid search on cluster
+nb_trials = 6   # (2 drop probs * 3 lrs)
+cluster.optimize_parallel_cluster_gpu(
+    YourMainFunction,
+    nb_trials=nb_trials,
+    job_name=hyperparams.experiment_name
+)
+```
+
+Running the above will launch 6 jobs, each with a different drop prob and learning rate combination.   
+The ```tunable``` parameter must be set to True to add that argument to the space of options, otherwise
+Test-Tube will use the ```default=value```.    
+
+
+## SLURM Flags   
+However you decide to submit your jobs, debugging requires a few flags. Without these flags, you'll
+see a nccl error instead of the actual error which caused the bug.   
+
+```sh
+export NCCL_DEBUG=INFO
+export PYTHONFAULTHANDLER=1
+```
+
+On some clusters you might need to set the network interface with this flag.   
+```sh
+export NCCL_SOCKET_IFNAME=^docker0,lo
+```   
+
+You might also need to load the latest version of NCCL  
+```sh
+module load NCCL/2.4.7-1-cuda.10.0
+```
+
+Finally, you must set the master port (usually a random number between 12k and 20k).   
+```sh
+# random port between 12k and 20k
+export MASTER_PORT=$((12000 + RANDOM % 20000))$   
+```
+
+## Simplest example.   
+1. Modify this script with your CoolModel file.   
+2. Update and submit [this bash script](https://github.com/williamFalcon/pytorch-lightning/blob/master/examples/new_project_templates/multi_node_examples/minimal_multi_node_demo_script.sh)   
+```bash
+squeue minimal_multi_node_demo_script.sh
+```
+
+## Grid search on a cluster   
+
+#### Option 1: Run on cluster using your own SLURM script    
+The trainer and model will work on a cluster if you configure your SLURM script correctly.   
+
+1. Update [this demo slurm script](https://github.com/williamFalcon/pytorch-lightning/blob/master/examples/new_project_templates/multi_node_examples/demo_script.sh).  
+2. Submit the script   
+```bash
+$ squeue demo_script.sh
+```
+
+Most people have some way they automatically generate their own scripts.  
+To run a grid search this way, you'd need a way to automatically generate scripts using all the combinations of 
+hyperparameters to search over.   
+
+#### Option 2: Use test-tube for SLURM script
+With test tube we can automatically generate slurm scripts for different hyperparameter options.   
+
+To run this demo:    
+```bash
+source activate YourCondaEnv
+
+python multi_node_cluster_auto_slurm.py --email your@email.com --gpu_partition your_partition --conda_env YourCondaEnv
+```
+
+That will submit 6 jobs. Each job will have a specific combination of hyperparams. Each job will also run on 2 nodes
+where each node has 8 gpus.   

examples/new_project_templates/multi_node_examples/demo_script.sh

@@ -0,0 +1,66 @@
+#!/bin/bash
+#
+# Auto-generated by test-tube (https://github.com/williamFalcon/test-tube)
+#################
+
+# set a job name
+#SBATCH --job-name=lightning_test
+#################
+
+# a file for job output, you can check job progress
+#SBATCH --output=/slurm_output_%j.out
+#################
+
+# a file for errors
+#SBATCH --error=/slurm_output_%j.err
+#################
+
+# time needed for job
+#SBATCH --time=01:00:00
+#################
+
+# gpus per node
+#SBATCH --gres=gpu:8
+#################
+
+# cpus per job
+#SBATCH --cpus-per-task=10
+#################
+
+# number of requested nodes
+#SBATCH --nodes=2
+#################
+
+# memory per node (0 means all)
+#SBATCH --mem=0
+#################
+
+# slurm will send a signal this far out before it kills the job
+#SBATCH --signal=USR1@300
+#################
+
+# comment
+#SBATCH --comment=lightning_demo
+#################
+
+# 1 task per gpu
+#SBATCH --ntasks-per-node=8
+#################
+
+source activate YourEnv
+
+# debugging flags (optional)
+export NCCL_DEBUG=INFO
+export PYTHONFAULTHANDLER=1
+
+# on your cluster you might need these:
+# set the network interface
+export NCCL_SOCKET_IFNAME=^docker0,lo
+
+# might need the latest cuda
+module load NCCL/2.4.7-1-cuda.10.0
+
+# random port between 12k and 20k
+export MASTER_PORT=$((12000 + RANDOM % 20000))$
+
+srun python multi_node_own_slurm_script.py

examples/new_project_templates/multi_node_examples/minimal_multi_node_demo.py

@@ -0,0 +1,24 @@
+from pytorch_lightning import Trainer
+from test_tube import Experiment
+import os
+
+
+def main():
+    # use the cool model from the main README.md
+    model = CoolModel()  # noqa: F821
+    exp = Experiment(save_dir=os.getcwd())
+
+    # train on 4 GPUs across 4 nodes
+    trainer = Trainer(
+        experiment=exp,
+        distributed_backend='ddp',
+        max_nb_epochs=10,
+        gpus=4,
+        nb_gpu_nodes=4
+    )
+
+    trainer.fit(model)
+
+
+if __name__ == '__main__':
+    main()

examples/new_project_templates/multi_node_examples/minimal_multi_node_demo_script.sh

@@ -0,0 +1,30 @@
+#!/bin/bash -l
+
+# SLURM SUBMIT SCRIPT
+#SBATCH --nodes=4
+#SBATCH --gres=gpu:4
+#SBATCH --ntasks-per-node=4
+#SBATCH --mem=0
+#SBATCH --time=0-02:00:00
+
+# activate conda env
+conda activate my_env
+
+# -------------------------
+# debugging flags (optional)
+# export NCCL_DEBUG=INFO
+# export PYTHONFAULTHANDLER=1
+
+# on your cluster you might need these:
+# set the network interface
+# export NCCL_SOCKET_IFNAME=^docker0,lo
+
+# might need the latest cuda
+# module load NCCL/2.4.7-1-cuda.10.0
+# -------------------------
+
+# random port between 12k and 20k
+export MASTER_PORT=$((12000 + RANDOM % 20000))
+
+# run script from above
+python minimal_multi_node_demo.py

examples/new_project_templates/multi_node_examples/multi_node_cluster_auto_slurm.py

@@ -7,11 +7,12 @@ from time import sleep
 import torch
 
 from test_tube import HyperOptArgumentParser, Experiment, SlurmCluster
-from pytorch_lightning.models.trainer import Trainer
+from pytorch_lightning import Trainer
 from pytorch_lightning.callbacks import EarlyStopping, ModelCheckpoint
 
 from examples.new_project_templates.lightning_module_template import LightningTemplateModel
 
+PORT = np.random.randint(12000, 20000, 1)[0]
 SEED = 2334
 torch.manual_seed(SEED)
 np.random.seed(SEED)
@@ -21,7 +22,7 @@ def main_local(hparams):
     main(hparams, None, None)
 
 
-def main(hparams, cluster, results_dict):
+def main(hparams, cluster):
     """
     Main training routine specific for this project
     :param hparams:
@@ -47,6 +48,7 @@ def main(hparams, cluster, results_dict):
         name=hyperparams.experiment_name,
         save_dir=hyperparams.test_tube_save_path,
         autosave=False,
+        version=hparams.hpc_exp_number,  # match the slurm job version number
         description='test demo'
     )
 
@@ -77,10 +79,9 @@ def main(hparams, cluster, results_dict):
     # ------------------------
     trainer = Trainer(
         experiment=exp,
-        cluster=cluster,
         checkpoint_callback=checkpoint,
         early_stop_callback=early_stop,
-        gpus=hparams.gpus,
+        gpus=hparams.per_experiment_nb_gpus,
         nb_gpu_nodes=hyperparams.nb_gpu_nodes
     )
 
@@ -99,7 +100,7 @@ def optimize_on_cluster(hyperparams):
     )
 
     # email for cluster coms
-    cluster.notify_job_status(email='add_email_here', on_done=True, on_fail=True)
+    cluster.notify_job_status(email=hyperparams.email, on_done=True, on_fail=True)
 
     # configure cluster
     cluster.per_experiment_nb_gpus = hyperparams.per_experiment_nb_gpus
@@ -109,7 +110,24 @@ def optimize_on_cluster(hyperparams):
     cluster.memory_mb_per_node = 0
 
     # any modules for code to run in env
-    cluster.add_command('source activate lightning')
+    cluster.add_command(f'source activate {hyperparams.conda_env}')
+
+    # set DDP master port
+    cluster.add_command(f'export MASTER_PORT={PORT}')
+
+    # OPTIONAL for debugging
+    # without these flags errors in your code will 
+    # appear to be nccl errors
+    cluster.add_command('export NCCL_DEBUG=INFO')
+    cluster.add_command('export PYTHONFAULTHANDLER=1')
+
+    # depending on your cluster config, you probably want
+    # to limit the wired connection device
+    # cluster.add_command('export NCCL_SOCKET_IFNAME=^docker0,lo')
+
+    # depending on your cluster, you might need to load
+    # the latest NCCL version
+    # cluster.load_modules(['NCCL/2.4.7-1-cuda.10.0'])
 
     # run only on 32GB voltas
     cluster.add_slurm_cmd(cmd='constraint', value='volta32gb',
@@ -121,7 +139,7 @@ def optimize_on_cluster(hyperparams):
     # creates and submits jobs to slurm
     cluster.optimize_parallel_cluster_gpu(
         main,
-        nb_trials=hyperparams.nb_hopt_trials,
+        nb_trials=hyperparams.num_hyperparam_trials,
         job_name=hyperparams.experiment_name
     )
 
@@ -139,15 +157,10 @@ if __name__ == '__main__':
     parent_parser = HyperOptArgumentParser(strategy='grid_search', add_help=False)
 
     # cluster args not defined inside the model
-    parent_parser.add_argument('--gpu_partition', type=str, help='consult your cluster manual')
 
-    # TODO: make 1 param
     parent_parser.add_argument('--per_experiment_nb_gpus', type=int,
-                               help='how many gpus to use in a node')
-    parent_parser.add_argument('--gpus', type=str, default='-1',
-                               help='how many gpus to use in the node')
-
-    parent_parser.add_argument('--nb_gpu_nodes', type=int, default=1,
+                               default=8, help='how many gpus to use in a node')
+    parent_parser.add_argument('--nb_gpu_nodes', type=int, default=2,
                                help='how many nodes to use in a cluster')
     parent_parser.add_argument('--test_tube_save_path', type=str, default=test_tube_dir,
                                help='where to save logs')
@@ -157,9 +170,15 @@ if __name__ == '__main__':
                                help='where to save model')
     parent_parser.add_argument('--experiment_name', type=str, default='pt_lightning_exp_a',
                                help='test tube exp name')
-    parent_parser.add_argument('--nb_hopt_trials', type=int, default=1,
+    parent_parser.add_argument('--num_hyperparam_trials', type=int, default=6,
                                help='how many grid search trials to run')
 
+    parent_parser.add_argument('--email', type=str, default='add@email.com',
+                               help='email for jobs')
+    parent_parser.add_argument('--conda_env', type=str, default='base',
+                               help='email for jobs')
+    parent_parser.add_argument('--gpu_partition', type=str, help='consult your cluster manual')
+
     # allow model to overwrite or extend args
     parser = LightningTemplateModel.add_model_specific_args(parent_parser, root_dir)
     hyperparams = parser.parse_args()

examples/new_project_templates/multi_node_examples/multi_node_own_slurm_script.py

@@ -0,0 +1,70 @@
+"""
+Multi-node example (GPU)
+"""
+import os
+import numpy as np
+import torch
+
+from test_tube import HyperOptArgumentParser, Experiment
+from pytorch_lightning import Trainer
+from examples.new_project_templates.lightning_module_template import LightningTemplateModel
+
+SEED = 2334
+torch.manual_seed(SEED)
+np.random.seed(SEED)
+
+
+def main(hparams):
+    """
+    Main training routine specific for this project
+    :param hparams:
+    :return:
+    """
+    # ------------------------
+    # 1 INIT LIGHTNING MODEL
+    # ------------------------
+    model = LightningTemplateModel(hparams)
+
+    # ------------------------
+    # 2 INIT TEST TUBE EXP
+    # ------------------------
+    # init experiment
+    exp = Experiment(
+        name='test_exp',
+        save_dir=hyperparams.log_dir,
+        autosave=False,
+        description='test demo'
+    )
+
+    # ------------------------
+    # 2 INIT TRAINER
+    # ------------------------
+    trainer = Trainer(
+        experiment=exp,
+        gpus=8,
+        nb_gpu_nodes=2
+    )
+
+    # ------------------------
+    # 5 START TRAINING
+    # ------------------------
+    trainer.fit(model)
+
+
+if __name__ == '__main__':
+    # use current dir for logging
+    root_dir = os.path.dirname(os.path.realpath(__file__))
+    log_dir = os.path.join(root_dir, 'pt_lightning_demo_logs')
+
+    parent_parser = HyperOptArgumentParser(strategy='grid_search', add_help=False)
+    parent_parser.add_argument('--log_dir', type=str, default=log_dir,
+                               help='where to save logs')
+
+    # allow model to overwrite or extend args
+    parser = LightningTemplateModel.add_model_specific_args(parent_parser, root_dir)
+    hyperparams = parser.parse_args()
+
+    # ---------------------
+    # RUN TRAINING
+    # ---------------------
+    main(hyperparams)

examples/new_project_templates/single_cpu_template.py

@@ -6,7 +6,7 @@ import numpy as np
 import torch
 
 from test_tube import HyperOptArgumentParser, Experiment
-from pytorch_lightning.models.trainer import Trainer
+from pytorch_lightning import Trainer
 from pytorch_lightning.callbacks import EarlyStopping, ModelCheckpoint
 
 from examples.new_project_templates.lightning_module_template import LightningTemplateModel
@@ -25,14 +25,11 @@ def main(hparams):
     # ------------------------
     # 1 INIT LIGHTNING MODEL
     # ------------------------
-    print('loading model...')
     model = LightningTemplateModel(hparams)
-    print('model built')
 
     # ------------------------
-    # 2 INIT TEST TUBE EXP
+    # 2 INIT EXP
     # ------------------------
-
     # init experiment
     exp = Experiment(
         name=hyperparams.experiment_name,

examples/new_project_templates/single_gpu_node_16bit_template.py

@@ -6,7 +6,7 @@ import numpy as np
 import torch
 
 from test_tube import HyperOptArgumentParser, Experiment
-from pytorch_lightning.models.trainer import Trainer
+from pytorch_lightning import Trainer
 from pytorch_lightning.callbacks import EarlyStopping, ModelCheckpoint
 
 from examples.new_project_templates.lightning_module_template import LightningTemplateModel

examples/new_project_templates/single_gpu_node_ddp_template.py

@@ -6,7 +6,7 @@ import numpy as np
 import torch
 
 from test_tube import HyperOptArgumentParser, Experiment
-from pytorch_lightning.models.trainer import Trainer
+from pytorch_lightning import Trainer
 from pytorch_lightning.callbacks import EarlyStopping, ModelCheckpoint
 
 from examples.new_project_templates.lightning_module_template import LightningTemplateModel

examples/new_project_templates/single_gpu_node_dp_template.py

@@ -6,7 +6,7 @@ import numpy as np
 import torch
 
 from test_tube import HyperOptArgumentParser, Experiment
-from pytorch_lightning.models.trainer import Trainer
+from pytorch_lightning import Trainer
 from pytorch_lightning.callbacks import EarlyStopping, ModelCheckpoint
 
 from examples.new_project_templates.lightning_module_template import LightningTemplateModel

examples/new_project_templates/trainer_cpu_template.py

@@ -2,7 +2,7 @@ import os
 import sys
 
 from test_tube import HyperOptArgumentParser, Experiment
-from pytorch_lightning.models.trainer import Trainer
+from pytorch_lightning import Trainer
 from pytorch_lightning.utilities.arg_parse import add_default_args
 from pytorch_lightning.callbacks.pt_callbacks import EarlyStopping, ModelCheckpoint
 

pytorch_lightning/__init__.py

@@ -1,4 +1,4 @@
-from .models.trainer import Trainer
+from .trainer.trainer import Trainer
 from .root_module.root_module import LightningModule
 from .root_module.decorators import data_loader
 

pytorch_lightning/callbacks/pt_callbacks.py

@@ -12,7 +12,6 @@ class Callback(object):
     # Properties
         params: dict. Training parameters
             (eg. verbosity, batch size, number of epochs...).
-        model: instance of `keras.models.Model`.
             Reference of the model being trained.
     The `logs` dictionary that callback methods
     take as argument will contain keys for quantities relevant to

pytorch_lightning/pt_overrides/override_data_parallel.py

@@ -56,6 +56,8 @@ class LightningDataParallel(DataParallel):
             # lightning
             if self.module.training:
                 return self.module.training_step(*inputs[0], **kwargs[0])
+            elif self.module.testing:
+                return self.module.test_step(*inputs[0], **kwargs[0])
             else:
                 return self.module.validation_step(*inputs[0], **kwargs[0])
 
@@ -89,6 +91,8 @@ class LightningDistributedDataParallel(DistributedDataParallel):
                 # lightning
                 if self.module.training:
                     output = self.module.training_step(*inputs[0], **kwargs[0])
+                elif self.module.testing:
+                    output = self.module.test_step(*inputs[0], **kwargs[0])
                 else:
                     output = self.module.validation_step(*inputs[0], **kwargs[0])
             else:
@@ -153,6 +157,10 @@ def parallel_apply(modules, inputs, kwargs_tup=None, devices=None):  # pragma: n
                 # CHANGE
                 if module.training:
                     output = module.training_step(*input, **kwargs)
+
+                elif module.testing:
+                    output = module.test_step(*input, **kwargs)
+
                 else:
                     output = module.validation_step(*input, **kwargs)
                 # ---------------

pytorch_lightning/root_module/hooks.py

@@ -5,7 +5,7 @@ class ModelHooks(torch.nn.Module):
 
     def on_sanity_check_start(self):
         """
-        Called before starting validate
+        Called before starting evaluate
         :return:
         """
         pass

pytorch_lightning/root_module/model_saving.py

@@ -1,12 +1,3 @@
-import os
-import re
-
-import torch
-
-from pytorch_lightning.pt_overrides.override_data_parallel import (
-    LightningDistributedDataParallel, LightningDataParallel)
-
-
 class ModelIO(object):
 
     def on_load_checkpoint(self, checkpoint):
@@ -41,221 +32,3 @@ class ModelIO(object):
         :return:
         """
         pass
-
-
-class TrainerIO(object):
-
-    def __get_model(self):
-        is_dp_module = isinstance(self.model, (LightningDistributedDataParallel,
-                                               LightningDataParallel))
-        model = self.model.module if is_dp_module else self.model
-        return model
-
-    # --------------------
-    # MODEL SAVE CHECKPOINT
-    # --------------------
-    def save_checkpoint(self, filepath):
-        checkpoint = self.dump_checkpoint()
-
-        # do the actual save
-        torch.save(checkpoint, filepath)
-
-    def restore(self, checkpoint_path, on_gpu):
-
-        if on_gpu:
-            checkpoint = torch.load(checkpoint_path)
-        else:
-            checkpoint = torch.load(checkpoint_path, map_location=lambda storage, loc: storage)
-
-        # load training state (affects trainer only)
-        self.restore_training_state(checkpoint)
-
-        # load model state
-        model = self.__get_model()
-
-        # load the state_dict on the model automatically
-        model.load_state_dict(checkpoint['state_dict'])
-
-    def dump_checkpoint(self):
-
-        checkpoint = {
-            'epoch': self.current_epoch,
-            'global_step': self.global_step
-        }
-
-        if self.checkpoint_callback is not None:
-            checkpoint['checkpoint_callback_best'] = self.checkpoint_callback.best
-
-        if self.early_stop_callback is not None:
-            checkpoint['early_stop_callback_wait'] = self.early_stop_callback.wait
-            checkpoint['early_stop_callback_patience'] = self.early_stop_callback.patience
-
-        # save optimizers
-        optimizer_states = []
-        for i, optimizer in enumerate(self.optimizers):
-            optimizer_states.append(optimizer.state_dict())
-
-        checkpoint['optimizer_states'] = optimizer_states
-
-        # save lr schedulers
-        lr_schedulers = []
-        for i, scheduler in enumerate(self.lr_schedulers):
-            lr_schedulers.append(scheduler.state_dict())
-
-        checkpoint['lr_schedulers'] = lr_schedulers
-
-        # add the state_dict from the model
-        model = self.__get_model()
-        checkpoint['state_dict'] = model.state_dict()
-
-        # give the model a chance to add a few things
-        model.on_save_checkpoint(checkpoint)
-
-        return checkpoint
-
-    # --------------------
-    # HPC IO
-    # --------------------
-    def enable_auto_hpc_walltime_manager(self):
-        if self.cluster is None:
-            return
-
-        # allow test tube to handle model check pointing automatically
-        # only if proc 0 so we don't trigger world_size resubmits
-        if self.proc_rank == 0:
-            self.cluster.set_checkpoint_save_function(
-                self.hpc_save,
-                kwargs={
-                    'folderpath': self.checkpoint_callback.filepath,
-                    'experiment': self.experiment
-                }
-            )
-
-        self.cluster.set_checkpoint_load_function(
-            self.hpc_load,
-            kwargs={
-                'folderpath': self.checkpoint_callback.filepath,
-                'on_gpu': self.on_gpu
-            }
-        )
-
-    def restore_training_state(self, checkpoint):
-        """
-        Restore trainer state.
-        Model will get its change to update
-        :param checkpoint:
-        :return:
-        """
-        if self.checkpoint_callback is not None:
-            self.checkpoint_callback.best = checkpoint['checkpoint_callback_best']
-
-        if self.early_stop_callback is not None:
-            self.early_stop_callback.wait = checkpoint['early_stop_callback_wait']
-            self.early_stop_callback.patience = checkpoint['early_stop_callback_patience']
-
-        self.global_step = checkpoint['global_step']
-        self.current_epoch = checkpoint['epoch']
-
-        # restore the optimizers
-        optimizer_states = checkpoint['optimizer_states']
-        for optimizer, opt_state in zip(self.optimizers, optimizer_states):
-            optimizer.load_state_dict(opt_state)
-
-        # restore the lr schedulers
-        lr_schedulers = checkpoint['lr_schedulers']
-        for scheduler, lrs_state in zip(self.lr_schedulers, lr_schedulers):
-            scheduler.load_state_dict(lrs_state)
-
-    # ----------------------------------
-    # PRIVATE OPS
-    # ----------------------------------
-    def hpc_save(self, folderpath, experiment):
-        # make sure the checkpoint folder exists
-        os.makedirs(folderpath, exist_ok=True)
-
-        # save exp to make sure we get all the metrics
-        experiment.save()
-
-        # close experiment to avoid issues
-        experiment.close()
-
-        ckpt_number = self.max_ckpt_in_folder(folderpath) + 1
-
-        if not os.path.exists(folderpath):
-            os.makedirs(folderpath, exist_ok=True)
-        filepath = '{}/hpc_ckpt_{}.ckpt'.format(folderpath, ckpt_number)
-
-        # give model a chance to do something on hpc_save
-        model = self.__get_model()
-        checkpoint = self.dump_checkpoint()
-
-        model.on_hpc_save(checkpoint)
-
-        # do the actual save
-        torch.save(checkpoint, filepath)
-
-        return filepath
-
-    def hpc_load(self, folderpath, on_gpu):
-        filepath = '{}/hpc_ckpt_{}.ckpt'.format(folderpath, self.max_ckpt_in_folder(folderpath))
-
-        if on_gpu:
-            checkpoint = torch.load(filepath)
-        else:
-            checkpoint = torch.load(filepath, map_location=lambda storage, loc: storage)
-
-        # load training state (affects trainer only)
-        self.restore_training_state(checkpoint)
-
-        # load model state
-        model = self.__get_model()
-
-        # load the state_dict on the model automatically
-        model.load_state_dict(checkpoint['state_dict'])
-
-        # call model hook
-        model.on_hpc_load(checkpoint)
-
-    def max_ckpt_in_folder(self, path, name_key='ckpt_'):
-        files = os.listdir(path)
-        files = [x for x in files if name_key in x]
-        if len(files) == 0:
-            return 0
-
-        ckpt_vs = []
-        for name in files:
-            name = name.split(name_key)[-1]
-            name = re.sub('[^0-9]', '', name)
-            ckpt_vs.append(int(name))
-
-        return max(ckpt_vs)
-
-
-def load_hparams_from_tags_csv(tags_csv):
-    from argparse import Namespace
-    import pandas as pd
-
-    tags_df = pd.read_csv(tags_csv)
-    dic = tags_df.to_dict(orient='records')
-
-    ns_dict = {row['key']: convert(row['value']) for row in dic}
-
-    ns = Namespace(**ns_dict)
-    return ns
-
-
-def convert(val):
-    constructors = [int, float, str]
-
-    if type(val) is str:
-        if val.lower() == 'true':
-            return True
-        if val.lower() == 'false':
-            return False
-
-    for c in constructors:
-        try:
-            return c(val)
-        except ValueError:
-            pass
-    return val

pytorch_lightning/root_module/root_module.py

@@ -2,7 +2,8 @@ import torch
 
 from pytorch_lightning.root_module.memory import ModelSummary
 from pytorch_lightning.root_module.grads import GradInformation
-from pytorch_lightning.root_module.model_saving import ModelIO, load_hparams_from_tags_csv
+from pytorch_lightning.trainer.trainer_io import load_hparams_from_tags_csv
+from pytorch_lightning.root_module.model_saving import ModelIO
 from pytorch_lightning.root_module.hooks import ModelHooks
 from pytorch_lightning.root_module.decorators import data_loader
 
@@ -55,6 +56,16 @@ class LightningModule(GradInformation, ModelIO, ModelHooks):
         """
         pass
 
+    def test_step(self, *args, **kwargs):
+        """
+        return whatever outputs will need to be aggregated in test_end
+        OPTIONAL
+        :param called with batch, batch_nb
+        additional: dataset_i if multiple val datasets used
+        :return:
+        """
+        pass
+
     def validation_end(self, outputs):
         """
         Outputs has the appended output after each validation step
@@ -64,6 +75,15 @@ class LightningModule(GradInformation, ModelIO, ModelHooks):
         """
         pass
 
+    def test_end(self, outputs):
+        """
+        Outputs has the appended output after each test step
+        OPTIONAL
+        :param outputs:
+        :return: dic_with_metrics for tqdm
+        """
+        pass
+
     def configure_optimizers(self):
         """
         Return a list of optimizers and a list of schedulers (could be empty)
@@ -110,7 +130,7 @@ class LightningModule(GradInformation, ModelIO, ModelHooks):
         return None
 
     @classmethod
-    def load_from_metrics(cls, weights_path, tags_csv, on_gpu, map_location=None):
+    def load_from_metrics(cls, weights_path, tags_csv, on_gpu):
         """
         Primary way of loading model from csv weights path
         :param weights_path:
@@ -122,13 +142,9 @@ class LightningModule(GradInformation, ModelIO, ModelHooks):
         hparams = load_hparams_from_tags_csv(tags_csv)
         hparams.__setattr__('on_gpu', on_gpu)
 
-        if on_gpu:
-            if map_location is not None:
-                checkpoint = torch.load(weights_path, map_location=map_location)
-            else:
-                checkpoint = torch.load(weights_path)
-        else:
-            checkpoint = torch.load(weights_path, map_location=lambda storage, loc: storage)
+        # load on CPU only to avoid OOM issues
+        # then its up to user to put back on GPUs
+        checkpoint = torch.load(weights_path, map_location=lambda storage, loc: storage)
 
         # load the state_dict on the model automatically
         model = cls(hparams)

pytorch_lightning/testing/__init__.py

@@ -1,3 +1,12 @@
 from .lm_test_module import LightningTestModel
-from .no_val_end_module import NoValEndTestModel
-from .no_val_module import NoValModel
+from .lm_test_module_base import LightningTestModelBase
+from .lm_test_module_mixins import (
+    LightningValidationStepMixin,
+    LightningValidationMixin,
+    LightningValidationStepMultipleDataloadersMixin,
+    LightningValidationMultipleDataloadersMixin,
+    LightningTestStepMixin,
+    LightningTestMixin,
+    LightningTestStepMultipleDataloadersMixin,
+    LightningTestMultipleDataloadersMixin,
+)

pytorch_lightning/testing/lm_test_module.py

@@ -14,262 +14,14 @@ from test_tube import HyperOptArgumentParser
 from pytorch_lightning.root_module.root_module import LightningModule
 from pytorch_lightning import data_loader
 
+from .lm_test_module_base import LightningTestModelBase
+from .lm_test_module_mixins import LightningValidationMixin, LightningTestMixin
 
-class LightningTestModel(LightningModule):
+
+class LightningTestModel(LightningValidationMixin, LightningTestMixin, LightningTestModelBase):
     """
-    Sample model to show how to define a template
+    Most common test case. Validation and test dataloaders
     """
 
-    def __init__(self, hparams, force_remove_distributed_sampler=False):
-        """
-        Pass in parsed HyperOptArgumentParser to the model
-        :param hparams:
-        """
-        # init superclass
-        super(LightningTestModel, self).__init__()
-        self.hparams = hparams
-
-        self.batch_size = hparams.batch_size
-
-        # if you specify an example input, the summary will show input/output for each layer
-        self.example_input_array = torch.rand(5, 28 * 28)
-
-        # remove to test warning for dist sampler
-        self.force_remove_distributed_sampler = force_remove_distributed_sampler
-
-        # build model
-        self.__build_model()
-
-    # ---------------------
-    # MODEL SETUP
-    # ---------------------
-    def __build_model(self):
-        """
-        Layout model
-        :return:
-        """
-        self.c_d1 = nn.Linear(in_features=self.hparams.in_features,
-                              out_features=self.hparams.hidden_dim)
-        self.c_d1_bn = nn.BatchNorm1d(self.hparams.hidden_dim)
-        self.c_d1_drop = nn.Dropout(self.hparams.drop_prob)
-
-        self.c_d2 = nn.Linear(in_features=self.hparams.hidden_dim,
-                              out_features=self.hparams.out_features)
-
-    # ---------------------
-    # TRAINING
-    # ---------------------
-    def forward(self, x):
-        """
-        No special modification required for lightning, define as you normally would
-        :param x:
-        :return:
-        """
-
-        x = self.c_d1(x)
-        x = torch.tanh(x)
-        x = self.c_d1_bn(x)
-        x = self.c_d1_drop(x)
-
-        x = self.c_d2(x)
-        logits = F.log_softmax(x, dim=1)
-
-        return logits
-
-    def loss(self, labels, logits):
-        nll = F.nll_loss(logits, labels)
-        return nll
-
-    def training_step(self, data_batch, batch_i):
-        """
-        Lightning calls this inside the training loop
-        :param data_batch:
-        :return:
-        """
-        # forward pass
-        x, y = data_batch
-        x = x.view(x.size(0), -1)
-
-        y_hat = self.forward(x)
-
-        # calculate loss
-        loss_val = self.loss(y, y_hat)
-
-        # in DP mode (default) make sure if result is scalar, there's another dim in the beginning
-        if self.trainer.use_dp:
-            loss_val = loss_val.unsqueeze(0)
-
-        # alternate possible outputs to test
-        if self.trainer.batch_nb % 1 == 0:
-            output = OrderedDict({
-                'loss': loss_val,
-                'prog': {'some_val': loss_val * loss_val}
-            })
-            return output
-        if self.trainer.batch_nb % 2 == 0:
-            return loss_val
-
-    def validation_step(self, data_batch, batch_i, dataloader_i):
-        """
-        Lightning calls this inside the validation loop
-        :param data_batch:
-        :return:
-        """
-        x, y = data_batch
-        x = x.view(x.size(0), -1)
-        y_hat = self.forward(x)
-
-        loss_val = self.loss(y, y_hat)
-
-        # acc
-        labels_hat = torch.argmax(y_hat, dim=1)
-        val_acc = torch.sum(y == labels_hat).item() / (len(y) * 1.0)
-        val_acc = torch.tensor(val_acc)
-
-        if self.on_gpu:
-            val_acc = val_acc.cuda(loss_val.device.index)
-
-        # in DP mode (default) make sure if result is scalar, there's another dim in the beginning
-        if self.trainer.use_dp:
-            loss_val = loss_val.unsqueeze(0)
-            val_acc = val_acc.unsqueeze(0)
-
-        # alternate possible outputs to test
-        if batch_i % 1 == 0:
-            output = OrderedDict({
-                'val_loss': loss_val,
-                'val_acc': val_acc,
-            })
-            return output
-        if batch_i % 2 == 0:
-            return val_acc
-
-        if batch_i % 3 == 0:
-            output = OrderedDict({
-                'val_loss': loss_val,
-                'val_acc': val_acc,
-                'test_dic': {'val_loss_a': loss_val}
-            })
-            return output
-        if batch_i % 5 == 0:
-            output = OrderedDict({
-                f'val_loss_{dataloader_i}': loss_val,
-                f'val_acc_{dataloader_i}': val_acc,
-            })
-            return output
-
-    def validation_end(self, outputs):
-        """
-        Called at the end of validation to aggregate outputs
-        :param outputs: list of individual outputs of each validation step
-        :return:
-        """
-        # if returned a scalar from validation_step, outputs is a list of tensor scalars
-        # we return just the average in this case (if we want)
-        # return torch.stack(outputs).mean()
-
-        val_loss_mean = 0
-        val_acc_mean = 0
-        for output in outputs:
-            val_loss_mean += output['val_loss']
-            val_acc_mean += output['val_acc']
-
-        val_loss_mean /= len(outputs)
-        val_acc_mean /= len(outputs)
-
-        tqdm_dic = {'val_loss': val_loss_mean.item(), 'val_acc': val_acc_mean.item()}
-        return tqdm_dic
-
     def on_tng_metrics(self, logs):
         logs['some_tensor_to_test'] = torch.rand(1)
-
-    # ---------------------
-    # TRAINING SETUP
-    # ---------------------
-    def configure_optimizers(self):
-        """
-        return whatever optimizers we want here
-        :return: list of optimizers
-        """
-        # try no scheduler for this model (testing purposes)
-        optimizer = optim.Adam(self.parameters(), lr=self.hparams.learning_rate)
-
-        # test returning only 1 list instead of 2
-        return optimizer
-
-    def __dataloader(self, train):
-        # init data generators
-        transform = transforms.Compose([transforms.ToTensor(),
-                                        transforms.Normalize((0.5,), (1.0,))])
-        dataset = MNIST(root=self.hparams.data_root, train=train,
-                        transform=transform, download=True)
-
-        # when using multi-node we need to add the datasampler
-        train_sampler = None
-        batch_size = self.hparams.batch_size
-
-        try:
-            if self.use_ddp and not self.force_remove_distributed_sampler:
-                train_sampler = DistributedSampler(dataset, rank=self.trainer.proc_rank)
-                batch_size = batch_size // self.trainer.world_size  # scale batch size
-        except Exception:
-            pass
-
-        should_shuffle = train_sampler is None
-        loader = DataLoader(
-            dataset=dataset,
-            batch_size=batch_size,
-            shuffle=should_shuffle,
-            sampler=train_sampler
-        )
-
-        return loader
-
-    @data_loader
-    def tng_dataloader(self):
-        return self.__dataloader(train=True)
-
-    @data_loader
-    def val_dataloader(self):
-        return [self.__dataloader(train=False), self.__dataloader(train=False)]
-
-    @data_loader
-    def test_dataloader(self):
-        return self.__dataloader(train=False)
-
-    @staticmethod
-    def add_model_specific_args(parent_parser, root_dir):  # pragma: no cover
-        """
-        Parameters you define here will be available to your model through self.hparams
-        :param parent_parser:
-        :param root_dir:
-        :return:
-        """
-        parser = HyperOptArgumentParser(strategy=parent_parser.strategy, parents=[parent_parser])
-
-        # param overwrites
-        # parser.set_defaults(gradient_clip=5.0)
-
-        # network params
-        parser.opt_list('--drop_prob', default=0.2, options=[0.2, 0.5], type=float, tunable=False)
-        parser.add_argument('--in_features', default=28 * 28, type=int)
-        parser.add_argument('--out_features', default=10, type=int)
-        # use 500 for CPU, 50000 for GPU to see speed difference
-        parser.add_argument('--hidden_dim', default=50000, type=int)
-
-        # data
-        parser.add_argument('--data_root', default=os.path.join(root_dir, 'mnist'), type=str)
-
-        # training params (opt)
-        parser.opt_list('--learning_rate', default=0.001 * 8, type=float,
-                        options=[0.0001, 0.0005, 0.001, 0.005],
-                        tunable=False)
-        parser.opt_list('--optimizer_name', default='adam', type=str,
-                        options=['adam'], tunable=False)
-
-        # if using 2 nodes with 4 gpus each the batch size here
-        #  (256) will be 256 / (2*8) = 16 per gpu
-        parser.opt_list('--batch_size', default=256 * 8, type=int,
-                        options=[32, 64, 128, 256], tunable=False,
-                        help='batch size will be divided over all gpus being used across all nodes')
-        return parser

pytorch_lightning/testing/lm_test_module_base.py

@@ -15,9 +15,10 @@ from pytorch_lightning.root_module.root_module import LightningModule
 from pytorch_lightning import data_loader
 
 
-class NoValModel(LightningModule):
+class LightningTestModelBase(LightningModule):
     """
-    Sample model to show how to define a template
+    Base LightningModule for testing. Implements only the required
+    interface
     """
 
     def __init__(self, hparams, force_remove_distributed_sampler=False):
@@ -26,7 +27,7 @@ class NoValModel(LightningModule):
         :param hparams:
         """
         # init superclass
-        super(NoValModel, self).__init__()
+        super(LightningTestModelBase, self).__init__()
         self.hparams = hparams
 
         self.batch_size = hparams.batch_size
@@ -109,9 +110,6 @@ class NoValModel(LightningModule):
         if self.trainer.batch_nb % 2 == 0:
             return loss_val
 
-    def on_tng_metrics(self, logs):
-        logs['some_tensor_to_test'] = torch.rand(1)
-
     # ---------------------
     # TRAINING SETUP
     # ---------------------
@@ -124,9 +122,9 @@ class NoValModel(LightningModule):
         optimizer = optim.Adam(self.parameters(), lr=self.hparams.learning_rate)
 
         # test returning only 1 list instead of 2
-        return [optimizer]
+        return optimizer
 
-    def __dataloader(self, train):
+    def _dataloader(self, train):
         # init data generators
         transform = transforms.Compose([transforms.ToTensor(),
                                         transforms.Normalize((0.5,), (1.0,))])
@@ -156,7 +154,7 @@ class NoValModel(LightningModule):
 
     @data_loader
     def tng_dataloader(self):
-        return self.__dataloader(train=True)
+        return self._dataloader(train=True)
 
     @staticmethod
     def add_model_specific_args(parent_parser, root_dir):  # pragma: no cover

pytorch_lightning/testing/lm_test_module_mixins.py

@@ -0,0 +1,387 @@
+import os
+from collections import OrderedDict
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch import optim
+from torch.utils.data import DataLoader
+from torch.utils.data.distributed import DistributedSampler
+from torchvision.datasets import MNIST
+from torchvision import transforms
+from test_tube import HyperOptArgumentParser
+
+from pytorch_lightning.root_module.root_module import LightningModule
+from pytorch_lightning import data_loader
+
+
+class LightningValidationStepMixin:
+    """
+    Add val_dataloader and validation_step methods for the case
+    when val_dataloader returns a single dataloader
+    """
+
+    @data_loader
+    def val_dataloader(self):
+        return self._dataloader(train=False)
+
+    def validation_step(self, data_batch, batch_i):
+        """
+        Lightning calls this inside the validation loop
+        :param data_batch:
+        :return:
+        """
+        x, y = data_batch
+        x = x.view(x.size(0), -1)
+        y_hat = self.forward(x)
+
+        loss_val = self.loss(y, y_hat)
+
+        # acc
+        labels_hat = torch.argmax(y_hat, dim=1)
+        val_acc = torch.sum(y == labels_hat).item() / (len(y) * 1.0)
+        val_acc = torch.tensor(val_acc)
+
+        if self.on_gpu:
+            val_acc = val_acc.cuda(loss_val.device.index)
+
+        # in DP mode (default) make sure if result is scalar, there's another dim in the beginning
+        if self.trainer.use_dp:
+            loss_val = loss_val.unsqueeze(0)
+            val_acc = val_acc.unsqueeze(0)
+
+        # alternate possible outputs to test
+        if batch_i % 1 == 0:
+            output = OrderedDict({
+                'val_loss': loss_val,
+                'val_acc': val_acc,
+            })
+            return output
+        if batch_i % 2 == 0:
+            return val_acc
+
+        if batch_i % 3 == 0:
+            output = OrderedDict({
+                'val_loss': loss_val,
+                'val_acc': val_acc,
+                'test_dic': {'val_loss_a': loss_val}
+            })
+            return output
+
+
+class LightningValidationMixin(LightningValidationStepMixin):
+    """
+    Add val_dataloader, validation_step, and validation_end methods for the case
+    when val_dataloader returns a single dataloader
+    """
+
+    def validation_end(self, outputs):
+        """
+        Called at the end of validation to aggregate outputs
+        :param outputs: list of individual outputs of each validation step
+        :return:
+        """
+        # if returned a scalar from validation_step, outputs is a list of tensor scalars
+        # we return just the average in this case (if we want)
+        # return torch.stack(outputs).mean()
+        val_loss_mean = 0
+        val_acc_mean = 0
+        for output in outputs:
+            val_loss = output['val_loss']
+
+            # reduce manually when using dp
+            if self.trainer.use_dp:
+                val_loss = torch.mean(val_loss)
+            val_loss_mean += val_loss
+
+            # reduce manually when using dp
+            val_acc = output['val_acc']
+            if self.trainer.use_dp:
+                val_acc = torch.mean(val_acc)
+
+            val_acc_mean += val_acc
+
+        val_loss_mean /= len(outputs)
+        val_acc_mean /= len(outputs)
+
+        tqdm_dic = {'val_loss': val_loss_mean.item(), 'val_acc': val_acc_mean.item()}
+        return tqdm_dic
+
+
+class LightningValidationStepMultipleDataloadersMixin:
+    """
+    Add val_dataloader and validation_step methods for the case
+    when val_dataloader returns multiple dataloaders
+    """
+
+    @data_loader
+    def val_dataloader(self):
+        return [self._dataloader(train=False), self._dataloader(train=False)]
+
+    def validation_step(self, data_batch, batch_i, dataloader_i):
+        """
+        Lightning calls this inside the validation loop
+        :param data_batch:
+        :return:
+        """
+        x, y = data_batch
+        x = x.view(x.size(0), -1)
+        y_hat = self.forward(x)
+
+        loss_val = self.loss(y, y_hat)
+
+        # acc
+        labels_hat = torch.argmax(y_hat, dim=1)
+        val_acc = torch.sum(y == labels_hat).item() / (len(y) * 1.0)
+        val_acc = torch.tensor(val_acc)
+
+        if self.on_gpu:
+            val_acc = val_acc.cuda(loss_val.device.index)
+
+        # in DP mode (default) make sure if result is scalar, there's another dim in the beginning
+        if self.trainer.use_dp:
+            loss_val = loss_val.unsqueeze(0)
+            val_acc = val_acc.unsqueeze(0)
+
+        # alternate possible outputs to test
+        if batch_i % 1 == 0:
+            output = OrderedDict({
+                'val_loss': loss_val,
+                'val_acc': val_acc,
+            })
+            return output
+        if batch_i % 2 == 0:
+            return val_acc
+
+        if batch_i % 3 == 0:
+            output = OrderedDict({
+                'val_loss': loss_val,
+                'val_acc': val_acc,
+                'test_dic': {'val_loss_a': loss_val}
+            })
+            return output
+        if batch_i % 5 == 0:
+            output = OrderedDict({
+                f'val_loss_{dataloader_i}': loss_val,
+                f'val_acc_{dataloader_i}': val_acc,
+            })
+            return output
+
+
+class LightningValidationMultipleDataloadersMixin(LightningValidationStepMultipleDataloadersMixin):
+    """
+    Add val_dataloader, validation_step, and validation_end methods for the case
+    when val_dataloader returns multiple dataloaders
+    """
+
+    def validation_end(self, outputs):
+        """
+        Called at the end of validation to aggregate outputs
+        :param outputs: list of individual outputs of each validation step
+        :return:
+        """
+        # if returned a scalar from validation_step, outputs is a list of tensor scalars
+        # we return just the average in this case (if we want)
+        # return torch.stack(outputs).mean()
+        val_loss_mean = 0
+        val_acc_mean = 0
+        i = 0
+        for dl_output in outputs:
+            for output in dl_output:
+                val_loss = output['val_loss']
+
+                # reduce manually when using dp
+                if self.trainer.use_dp:
+                    val_loss = torch.mean(val_loss)
+                val_loss_mean += val_loss
+
+                # reduce manually when using dp
+                val_acc = output['val_acc']
+                if self.trainer.use_dp:
+                    val_acc = torch.mean(val_acc)
+
+                val_acc_mean += val_acc
+                i += 1
+
+        val_loss_mean /= i
+        val_acc_mean /= i
+
+        tqdm_dic = {'val_loss': val_loss_mean.item(), 'val_acc': val_acc_mean.item()}
+        return tqdm_dic
+
+
+class LightningTestStepMixin:
+
+    @data_loader
+    def test_dataloader(self):
+        return self._dataloader(train=False)
+
+    def test_step(self, data_batch, batch_i):
+        """
+        Lightning calls this inside the validation loop
+        :param data_batch:
+        :return:
+        """
+        x, y = data_batch
+        x = x.view(x.size(0), -1)
+        y_hat = self.forward(x)
+
+        loss_test = self.loss(y, y_hat)
+
+        # acc
+        labels_hat = torch.argmax(y_hat, dim=1)
+        test_acc = torch.sum(y == labels_hat).item() / (len(y) * 1.0)
+        test_acc = torch.tensor(test_acc)
+
+        if self.on_gpu:
+            test_acc = test_acc.cuda(loss_test.device.index)
+
+        # in DP mode (default) make sure if result is scalar, there's another dim in the beginning
+        if self.trainer.use_dp:
+            loss_test = loss_test.unsqueeze(0)
+            test_acc = test_acc.unsqueeze(0)
+
+        # alternate possible outputs to test
+        if batch_i % 1 == 0:
+            output = OrderedDict({
+                'test_loss': loss_test,
+                'test_acc': test_acc,
+            })
+            return output
+        if batch_i % 2 == 0:
+            return test_acc
+
+        if batch_i % 3 == 0:
+            output = OrderedDict({
+                'test_loss': loss_test,
+                'test_acc': test_acc,
+                'test_dic': {'test_loss_a': loss_test}
+            })
+            return output
+
+
+class LightningTestMixin(LightningTestStepMixin):
+    def test_end(self, outputs):
+        """
+        Called at the end of validation to aggregate outputs
+        :param outputs: list of individual outputs of each validation step
+        :return:
+        """
+        # if returned a scalar from test_step, outputs is a list of tensor scalars
+        # we return just the average in this case (if we want)
+        # return torch.stack(outputs).mean()
+        test_loss_mean = 0
+        test_acc_mean = 0
+        for output in outputs:
+            test_loss = output['test_loss']
+
+            # reduce manually when using dp
+            if self.trainer.use_dp:
+                test_loss = torch.mean(test_loss)
+            test_loss_mean += test_loss
+
+            # reduce manually when using dp
+            test_acc = output['test_acc']
+            if self.trainer.use_dp:
+                test_acc = torch.mean(test_acc)
+
+            test_acc_mean += test_acc
+
+        test_loss_mean /= len(outputs)
+        test_acc_mean /= len(outputs)
+
+        tqdm_dic = {'test_loss': test_loss_mean.item(), 'test_acc': test_acc_mean.item()}
+        return tqdm_dic
+
+
+class LightningTestStepMultipleDataloadersMixin:
+
+    @data_loader
+    def test_dataloader(self):
+        return [self._dataloader(train=False), self._dataloader(train=False)]
+
+    def test_step(self, data_batch, batch_i, dataloader_i):
+        """
+        Lightning calls this inside the validation loop
+        :param data_batch:
+        :return:
+        """
+        x, y = data_batch
+        x = x.view(x.size(0), -1)
+        y_hat = self.forward(x)
+
+        loss_test = self.loss(y, y_hat)
+
+        # acc
+        labels_hat = torch.argmax(y_hat, dim=1)
+        test_acc = torch.sum(y == labels_hat).item() / (len(y) * 1.0)
+        test_acc = torch.tensor(test_acc)
+
+        if self.on_gpu:
+            test_acc = test_acc.cuda(loss_test.device.index)
+
+        # in DP mode (default) make sure if result is scalar, there's another dim in the beginning
+        if self.trainer.use_dp:
+            loss_test = loss_test.unsqueeze(0)
+            test_acc = test_acc.unsqueeze(0)
+
+        # alternate possible outputs to test
+        if batch_i % 1 == 0:
+            output = OrderedDict({
+                'test_loss': loss_test,
+                'test_acc': test_acc,
+            })
+            return output
+        if batch_i % 2 == 0:
+            return test_acc
+
+        if batch_i % 3 == 0:
+            output = OrderedDict({
+                'test_loss': loss_test,
+                'test_acc': test_acc,
+                'test_dic': {'test_loss_a': loss_test}
+            })
+            return output
+        if batch_i % 5 == 0:
+            output = OrderedDict({
+                f'test_loss_{dataloader_i}': loss_test,
+                f'test_acc_{dataloader_i}': test_acc,
+            })
+            return output
+
+
+class LightningTestMultipleDataloadersMixin(LightningTestStepMultipleDataloadersMixin):
+    def test_end(self, outputs):
+        """
+        Called at the end of validation to aggregate outputs
+        :param outputs: list of individual outputs of each validation step
+        :return:
+        """
+        # if returned a scalar from test_step, outputs is a list of tensor scalars
+        # we return just the average in this case (if we want)
+        # return torch.stack(outputs).mean()
+        test_loss_mean = 0
+        test_acc_mean = 0
+        i = 0
+        for dl_output in outputs:
+            for output in dl_output:
+                test_loss = output['test_loss']
+
+                # reduce manually when using dp
+                if self.trainer.use_dp:
+                    test_loss = torch.mean(test_loss)
+                test_loss_mean += test_loss
+
+                # reduce manually when using dp
+                test_acc = output['test_acc']
+                if self.trainer.use_dp:
+                    test_acc = torch.mean(test_acc)
+
+                test_acc_mean += test_acc
+                i += 1
+
+        test_loss_mean /= i
+        test_acc_mean /= i
+
+        tqdm_dic = {'test_loss': test_loss_mean.item(), 'test_acc': test_acc_mean.item()}
+        return tqdm_dic

pytorch_lightning/testing/no_val_end_module.py

@@ -1,247 +0,0 @@
-import os
-from collections import OrderedDict
-
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-from torch import optim
-from torch.utils.data import DataLoader
-from torch.utils.data.distributed import DistributedSampler
-from torchvision.datasets import MNIST
-from torchvision import transforms
-from test_tube import HyperOptArgumentParser
-
-from pytorch_lightning.root_module.root_module import LightningModule
-from pytorch_lightning import data_loader
-
-
-class NoValEndTestModel(LightningModule):
-    """
-    Sample model to show how to define a template
-    """
-
-    def __init__(self, hparams, force_remove_distributed_sampler=False):
-        """
-        Pass in parsed HyperOptArgumentParser to the model
-        :param hparams:
-        """
-        # init superclass
-        super(NoValEndTestModel, self).__init__()
-        self.hparams = hparams
-
-        self.batch_size = hparams.batch_size
-
-        # if you specify an example input, the summary will show input/output for each layer
-        self.example_input_array = torch.rand(5, 28 * 28)
-
-        # remove to test warning for dist sampler
-        self.force_remove_distributed_sampler = force_remove_distributed_sampler
-
-        # build model
-        self.__build_model()
-
-    # ---------------------
-    # MODEL SETUP
-    # ---------------------
-    def __build_model(self):
-        """
-        Layout model
-        :return:
-        """
-        self.c_d1 = nn.Linear(in_features=self.hparams.in_features,
-                              out_features=self.hparams.hidden_dim)
-        self.c_d1_bn = nn.BatchNorm1d(self.hparams.hidden_dim)
-        self.c_d1_drop = nn.Dropout(self.hparams.drop_prob)
-
-        self.c_d2 = nn.Linear(in_features=self.hparams.hidden_dim,
-                              out_features=self.hparams.out_features)
-
-    # ---------------------
-    # TRAINING
-    # ---------------------
-    def forward(self, x):
-        """
-        No special modification required for lightning, define as you normally would
-        :param x:
-        :return:
-        """
-
-        x = self.c_d1(x)
-        x = torch.tanh(x)
-        x = self.c_d1_bn(x)
-        x = self.c_d1_drop(x)
-
-        x = self.c_d2(x)
-        logits = F.log_softmax(x, dim=1)
-
-        return logits
-
-    def loss(self, labels, logits):
-        nll = F.nll_loss(logits, labels)
-        return nll
-
-    def training_step(self, data_batch, batch_i):
-        """
-        Lightning calls this inside the training loop
-        :param data_batch:
-        :return:
-        """
-        # forward pass
-        x, y = data_batch
-        x = x.view(x.size(0), -1)
-
-        y_hat = self.forward(x)
-
-        # calculate loss
-        loss_val = self.loss(y, y_hat)
-
-        # in DP mode (default) make sure if result is scalar, there's another dim in the beginning
-        if self.trainer.use_dp:
-            loss_val = loss_val.unsqueeze(0)
-
-        # alternate possible outputs to test
-        if self.trainer.batch_nb % 1 == 0:
-            output = OrderedDict({
-                'loss': loss_val,
-                'prog': {'some_val': loss_val * loss_val}
-            })
-            return output
-        if self.trainer.batch_nb % 2 == 0:
-            return loss_val
-
-    def validation_step(self, data_batch, batch_nb):
-        """
-        Lightning calls this inside the validation loop
-        :param data_batch:
-        :return:
-        """
-        x, y = data_batch
-        x = x.view(x.size(0), -1)
-        y_hat = self.forward(x)
-
-        loss_val = self.loss(y, y_hat)
-
-        # acc
-        labels_hat = torch.argmax(y_hat, dim=1)
-        val_acc = torch.sum(y == labels_hat).item() / (len(y) * 1.0)
-        val_acc = torch.tensor(val_acc)
-
-        if self.on_gpu:
-            val_acc = val_acc.cuda(loss_val.device.index)
-
-        # in DP mode (default) make sure if result is scalar, there's another dim in the beginning
-        if self.trainer.use_dp:
-            loss_val = loss_val.unsqueeze(0)
-            val_acc = val_acc.unsqueeze(0)
-
-        # alternate possible outputs to test
-        if batch_nb % 1 == 0:
-            output = OrderedDict({
-                'val_loss': loss_val,
-                'val_acc': val_acc,
-            })
-            return output
-        if batch_nb % 2 == 0:
-            return val_acc
-
-        if batch_nb % 3 == 0:
-            output = OrderedDict({
-                'val_loss': loss_val,
-                'val_acc': val_acc,
-                'test_dic': {'val_loss_a': loss_val}
-            })
-            return output
-
-    def on_tng_metrics(self, logs):
-        logs['some_tensor_to_test'] = torch.rand(1)
-
-    # ---------------------
-    # TRAINING SETUP
-    # ---------------------
-    def configure_optimizers(self):
-        """
-        return whatever optimizers we want here
-        :return: list of optimizers
-        """
-        # try no scheduler for this model (testing purposes)
-        optimizer = optim.Adam(self.parameters(), lr=self.hparams.learning_rate)
-
-        # test returning only 1 list instead of 2
-        return [optimizer]
-
-    def __dataloader(self, train):
-        # init data generators
-        transform = transforms.Compose([transforms.ToTensor(),
-                                        transforms.Normalize((0.5,), (1.0,))])
-        dataset = MNIST(root=self.hparams.data_root, train=train,
-                        transform=transform, download=True)
-
-        # when using multi-node we need to add the datasampler
-        train_sampler = None
-        batch_size = self.hparams.batch_size
-
-        try:
-            if self.use_ddp and not self.force_remove_distributed_sampler:
-                train_sampler = DistributedSampler(dataset, rank=self.trainer.proc_rank)
-                batch_size = batch_size // self.trainer.world_size  # scale batch size
-        except Exception:
-            pass
-
-        should_shuffle = train_sampler is None
-        loader = DataLoader(
-            dataset=dataset,
-            batch_size=batch_size,
-            shuffle=should_shuffle,
-            sampler=train_sampler
-        )
-
-        return loader
-
-    @data_loader
-    def tng_dataloader(self):
-        return self.__dataloader(train=True)
-
-    @data_loader
-    def val_dataloader(self):
-        return self.__dataloader(train=False)
-
-    @data_loader
-    def test_dataloader(self):
-        return self.__dataloader(train=False)
-
-    @staticmethod
-    def add_model_specific_args(parent_parser, root_dir):  # pragma: no cover
-        """
-        Parameters you define here will be available to your model through self.hparams
-        :param parent_parser:
-        :param root_dir:
-        :return:
-        """
-        parser = HyperOptArgumentParser(strategy=parent_parser.strategy, parents=[parent_parser])
-
-        # param overwrites
-        # parser.set_defaults(gradient_clip=5.0)
-
-        # network params
-        parser.opt_list('--drop_prob', default=0.2, options=[0.2, 0.5], type=float, tunable=False)
-        parser.add_argument('--in_features', default=28 * 28, type=int)
-        parser.add_argument('--out_features', default=10, type=int)
-        # use 500 for CPU, 50000 for GPU to see speed difference
-        parser.add_argument('--hidden_dim', default=50000, type=int)
-
-        # data
-        parser.add_argument('--data_root', default=os.path.join(root_dir, 'mnist'), type=str)
-
-        # training params (opt)
-        parser.opt_list('--learning_rate', default=0.001 * 8, type=float,
-                        options=[0.0001, 0.0005, 0.001, 0.005],
-                        tunable=False)
-        parser.opt_list('--optimizer_name', default='adam', type=str,
-                        options=['adam'], tunable=False)
-
-        # if using 2 nodes with 4 gpus each the batch size here
-        #  (256) will be 256 / (2*8) = 16 per gpu
-        parser.opt_list('--batch_size', default=256 * 8, type=int,
-                        options=[32, 64, 128, 256], tunable=False,
-                        help='batch size will be divided over all gpus being used across all nodes')
-        return parser

pytorch_lightning/trainer/trainer.py

@@ -14,12 +14,14 @@ import torch.multiprocessing as mp
 import torch.distributed as dist
 from torch.optim.optimizer import Optimizer
 
+from pytorch_lightning.root_module.root_module import LightningModule
 from pytorch_lightning.root_module.memory import get_gpu_memory_map
-from pytorch_lightning.root_module.model_saving import TrainerIO
+from pytorch_lightning.trainer.trainer_io import TrainerIO
 from pytorch_lightning.pt_overrides.override_data_parallel import (
     LightningDistributedDataParallel, LightningDataParallel)
 from pytorch_lightning.callbacks import GradientAccumulationScheduler
 from pytorch_lightning.utilities.debugging import MisconfigurationException
+import pdb
 
 try:
     from apex import amp
@@ -56,11 +58,10 @@ class Trainer(TrainerIO):
                  early_stop_callback=None,
                  checkpoint_callback=None,
                  gradient_clip=0,
-                 cluster=None,
                  process_position=0,
-                 current_gpu_name=0,
                  nb_gpu_nodes=1,
                  gpus=None,
+                 log_gpu_memory=False,
                  show_progress_bar=True,
                  overfit_pct=0.0,
                  track_grad_norm=-1,
@@ -75,88 +76,60 @@ class Trainer(TrainerIO):
                  val_check_interval=1.0,
                  log_save_interval=100,
                  add_log_row_interval=10,
-                 distributed_backend='dp',
+                 distributed_backend=None,
                  use_amp=False,
                  print_nan_grads=False,
                  print_weights_summary=True,
+                 weights_save_path=None,
                  amp_level='O2',
                  nb_sanity_val_steps=5):
         """
 
         :param experiment: Test-tube experiment
-        :param early_stop_callback: from pytorch_lightning import EarlyStopping
-        :param checkpoint_callback: from pytorch_lightning import Checkpoint
-        :param gradient_clip:
-        :param cluster:
-        :param process_position:
-        :param current_gpu_name:
-        :param nb_gpu_nodes:
-        :param gpus:
-        :param show_progress_bar:
-        :param overfit_pct:
-        :param track_grad_norm:
-        :param check_val_every_n_epoch:
-        :param fast_dev_run:
-        :param accumulate_grad_batches:
-        :param max_nb_epochs:
-        :param min_nb_epochs:
-        :param train_percent_check:
-        :param val_percent_check:
-        :param test_percent_check:
-        :param val_check_interval:
-        :param log_save_interval:
-        :param add_log_row_interval:
-        :param distributed_backend:
-            'do' to use DistributedParallel, 'dp' to use DistributedDataParallel, 'n' to use none
-        :param use_amp:
-        :param print_nan_grads:
-        :param print_weights_summary:
-        :param amp_level:
-        :param nb_sanity_val_steps:
+        :param early_stop_callback: Callback for early stopping
+        :param checkpoint_callback: Callback for checkpointing
+        :param gradient_clip: int. 0 means don't clip.
+        :param process_position: shown in the tqdm bar
+        :param nb_gpu_nodes: number of GPU nodes
+        :param gpus: int. (ie: 2 gpus) OR list to specify which GPUs [0, 1] or '0,1'
+        :param log_gpu_memory: Bool. If true, adds memory logs
+        :param show_progress_bar: Bool. If true shows tqdm bar
+        :param overfit_pct: float. uses this much of all datasets
+        :param track_grad_norm: int. -1 no tracking. Otherwise tracks that norm
+        :param check_val_every_n_epoch: int. check val every n train epochs
+        :param fast_dev_run: Bool. runs full iteration over everything to find bugs
+        :param accumulate_grad_batches: int. Accumulates grads every k batches
+        :param max_nb_epochs: int.
+        :param min_nb_epochs: int.
+        :param train_percent_check: int. How much of train set to check
+        :param val_percent_check: int. How much of val set to check
+        :param test_percent_check: int. How much of test set to check
+        :param val_check_interval: int. Check val this frequently within a train epoch
+        :param log_save_interval: int. Writes logs to disk this often
+        :param add_log_row_interval: int. How often to add logging rows
+        :param distributed_backend: str. dp, or ddp.
+        :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_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.
         """
         # Transfer params
         self.nb_gpu_nodes = nb_gpu_nodes
+        self.log_gpu_memory = log_gpu_memory
         self.gradient_clip = gradient_clip
         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.experiment = experiment
-        self.exp_save_path = None
-        if self.experiment is not None:
-            self.exp_save_path = experiment.get_data_path(experiment.name, experiment.version)
-        self.cluster = cluster
         self.process_position = process_position
-        self.current_gpu_name = current_gpu_name
         self.print_weights_summary = print_weights_summary
-        self.checkpoint_callback = checkpoint_callback
-
-        if self.checkpoint_callback is not None:
-            self.checkpoint_callback.save_function = self.save_checkpoint
-
-        self.early_stop = early_stop_callback
-        self.model = None
         self.max_nb_epochs = max_nb_epochs
-        if isinstance(accumulate_grad_batches, dict):
-            self.accumulation_scheduler = GradientAccumulationScheduler(accumulate_grad_batches)
-        elif isinstance(accumulate_grad_batches, int):
-            schedule = {1: accumulate_grad_batches}
-            self.accumulation_scheduler = GradientAccumulationScheduler(schedule)
-        else:
-            raise TypeError("Gradient accumulation supports only int and dict types")
-        self.early_stop_callback = early_stop_callback
         self.min_nb_epochs = min_nb_epochs
         self.nb_sanity_val_steps = nb_sanity_val_steps
-        self.lr_schedulers = []
-        self.amp_level = amp_level
         self.print_nan_grads = print_nan_grads
-        self.data_parallel_device_ids = None
-        self.world_size = 1
-        self.node_rank = 0
-        self.use_ddp = False
-        self.use_dp = False
-        self.single_gpu = False
 
         # training bookeeping
         self.total_batch_nb = 0
@@ -167,71 +140,53 @@ class Trainer(TrainerIO):
         self.nb_val_batches = 0
         self.nb_tng_batches = 0
         self.nb_test_batches = 0
-
-        # gpus come in as a string.
-        # if gpus = -1 then use all available devices
-        # otherwise, split the string using commas
-        if gpus is not None:
-            if type(gpus) is list:
-                self.data_parallel_device_ids = gpus
-            elif type(gpus) is str:
-                if gpus == '-1':
-                    self.data_parallel_device_ids = list(range(0, torch.cuda.device_count()))
-                else:
-                    self.data_parallel_device_ids = [int(x.strip()) for x in gpus.split(',')]
-            else:
-                raise Exception('gpus has to be a string or list of ids')
-
-            # set the correct cuda visible devices (using pci order)
-            os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
-            os.environ["CUDA_VISIBLE_DEVICES"] = ','.join([str(x) for x in
-                                                           self.data_parallel_device_ids])
-            print('VISIBLE GPUS: %r' % os.environ["CUDA_VISIBLE_DEVICES"])
-
-        # make DP and DDP mutually exclusive
-        # single GPU will also use DP with devices=[0]
-        requested_gpus = self.data_parallel_device_ids is not None
-        if requested_gpus and len(self.data_parallel_device_ids) > 0:
-            self.use_dp = distributed_backend == 'dp'
-            self.use_ddp = distributed_backend == 'ddp'
-
-            # use ddp automatically if nb_gpu_nodes > 1
-            if nb_gpu_nodes > 1 and self.use_dp:  # pragma: no cover
-                self.use_ddp = True
-                self.use_dp = False
-                w = 'DataParallel does not support nb_gpu_nodes > 1. ' \
-                    'Switching to DistributedDataParallel for you. ' \
-                    'To silence this warning set distributed_backend=ddp'
-                warnings.warn(w)
-
-        # remove dp and ddp when requesting single gpu
-        if self.data_parallel_device_ids is not None and len(self.data_parallel_device_ids) == 1:
-            self.use_ddp = False
-            self.use_dp = False
-            self.single_gpu = True
-
-        # extract SLURM flag vars
-        # whenever we have the correct number of tasks, we let slurm manage processes
-        # otherwise we launch the required number of processes
-        if self.use_ddp:
-            self.nb_requested_gpus = len(self.data_parallel_device_ids) * self.nb_gpu_nodes
-            self.nb_slurm_tasks = 0
-            try:
-                self.nb_slurm_tasks = int(os.environ['SLURM_NTASKS'])
-                self.is_slurm_managing_tasks = self.nb_slurm_tasks == self.nb_requested_gpus
-            except Exception:
-                # likely not on slurm, so set the slurm managed flag to false
-                self.is_slurm_managing_tasks = False
-
-        # process info
-        self.proc_rank = 0
+        self.tng_dataloader = None
+        self.test_dataloader = None
+        self.val_dataloader = None
 
         # training state
+        self.model = None
+        self.testing = False
+        self.lr_schedulers = []
         self.optimizers = None
         self.global_step = 0
         self.current_epoch = 0
         self.total_batches = 0
 
+        # configure early stop callback
+        self.early_stop_callback = early_stop_callback
+
+        # configure weights save path
+        self.__configure_weights_path(checkpoint_callback, weights_save_path)
+
+        # configure experiment
+        self.experiment = experiment
+        self.exp_save_path = None
+        if self.experiment is not None:
+            self.exp_save_path = experiment.get_data_path(experiment.name, experiment.version)
+
+        # accumulated grads
+        self.__configure_accumulated_gradients(accumulate_grad_batches)
+
+        # allow int, string and gpu list
+        self.data_parallel_device_ids = self.__parse_gpu_ids(gpus)
+        self.root_gpu = self.__set_root_gpu(self.data_parallel_device_ids)
+
+        # distributed backend choice
+        self.use_ddp = False
+        self.use_dp = False
+        self.single_gpu = False
+        self.__set_distributed_mode(distributed_backend, nb_gpu_nodes)
+
+        # init flags for SLURM+ddp to work
+        self.proc_rank = 0
+        self.world_size = 1
+        self.node_rank = 0
+        self.__configure_slurm_ddp(nb_gpu_nodes)
+
+        # nvidia setup
+        self.__set_nvidia_flags(self.is_slurm_managing_tasks, self.data_parallel_device_ids)
+
         # can't init progress bar here because starting a new process
         # means the prog_bar won't survive pickling
         self.show_progress_bar = show_progress_bar
@@ -241,17 +196,36 @@ class Trainer(TrainerIO):
         self.val_check_interval = val_check_interval
         self.add_log_row_interval = add_log_row_interval
 
-        # dataloaders
-        self.tng_dataloader = None
-        self.test_dataloader = None
-        self.val_dataloader = None
-
         # how much of the data to use
         self.__determine_data_use_amount(train_percent_check, val_percent_check,
                                          test_percent_check, overfit_pct)
-        print('gpu available: {}, used: {}'.format(torch.cuda.is_available(), self.on_gpu))
 
         # 16 bit mixed precision training using apex
+        self.amp_level = amp_level
+        self.__init_amp(use_amp)
+
+    def __configure_weights_path(self, checkpoint_callback, weights_save_path):
+        """
+        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
+
+        # configure checkpoint callback
+        self.checkpoint_callback = checkpoint_callback
+        if self.checkpoint_callback is not None:
+            self.checkpoint_callback.save_function = self.save_checkpoint
+
+            # 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 workingdir
+        if self.weights_save_path is None:
+            self.weights_save_path = os.getcwd()
+
+    def __init_amp(self, use_amp):
         self.use_amp = use_amp and APEX_AVAILABLE
         if self.use_amp:
             print('using 16bit precision')
@@ -266,36 +240,136 @@ class Trainer(TrainerIO):
             """
             raise ModuleNotFoundError(msg)
 
-    def restore_state_if_existing_checkpoint(self):
-        # restore trainer state and model if there is a weight for this experiment
-        last_epoch = -1
-        last_ckpt_name = None
+    def __configure_accumulated_gradients(self, accumulate_grad_batches):
+        if isinstance(accumulate_grad_batches, dict):
+            self.accumulation_scheduler = GradientAccumulationScheduler(accumulate_grad_batches)
+        elif isinstance(accumulate_grad_batches, int):
+            schedule = {1: accumulate_grad_batches}
+            self.accumulation_scheduler = GradientAccumulationScheduler(schedule)
+        else:
+            raise TypeError("Gradient accumulation supports only int and dict types")
 
-        # do nothing if there's not dir or callback
-        no_ckpt_callback = self.checkpoint_callback is None
-        if no_ckpt_callback or not os.path.exists(self.checkpoint_callback.filepath):
+    def __parse_gpu_ids(self, gpus):
+        """
+        :param gpus: Int, string or list of ids
+        :return:
+        """
+        # if gpus = -1 then use all available devices
+        # otherwise, split the string using commas
+        if gpus is not None:
+            if type(gpus) is list:
+                gpus = gpus
+            elif type(gpus) is str:
+                if gpus == '-1':
+                    gpus = list(range(0, torch.cuda.device_count()))
+                else:
+                    gpus = [int(x.strip()) for x in gpus.split(',')]
+            elif type(gpus) is int:
+                gpus = gpus
+            else:
+                raise Exception('gpus has to be a string, int or list of ints')
+
+        return gpus
+
+    def __set_root_gpu(self, gpus):
+        if gpus is None:
+            return None
+
+        # set root gpu
+        root_gpu = 0
+        if type(gpus) is list:
+            root_gpu = gpus[0]
+
+        return root_gpu
+
+    @property
+    def num_gpus(self):
+        gpus = self.data_parallel_device_ids
+        if gpus is None:
+            return 0
+        if type(gpus) is list:
+            return len(gpus)
+        if type(gpus) is int:
+            return gpus
+
+        m = 'gpus must be int, none or list of ints'
+        raise MisconfigurationException(m)
+
+    def __set_distributed_mode(self, distributed_backend, nb_gpu_nodes):
+        # skip for CPU
+        if self.num_gpus == 0:
             return
 
-        # find last epoch
-        checkpoints = os.listdir(self.checkpoint_callback.filepath)
-        for name in checkpoints:
-            # ignore hpc ckpts
-            if 'hpc_' in name:
-                continue
+        # single GPU case
+        if self.num_gpus == 1:
+            self.single_gpu = True
 
-            if '.ckpt' in name:
-                epoch = name.split('epoch_')[1]
-                epoch = int(re.sub('[^0-9]', '', epoch))
+            if distributed_backend is not None:
+                self.use_dp = distributed_backend == 'dp'
+                self.use_ddp = distributed_backend == 'ddp'
 
-                if epoch > last_epoch:
-                    last_epoch = epoch
-                    last_ckpt_name = name
+        # multiple GPU case
+        elif self.num_gpus > 1:
+            if distributed_backend is not None:
+                # DP, DDP case
+                self.use_dp = distributed_backend == 'dp'
+                self.use_ddp = distributed_backend == 'ddp'
 
-        # restore last checkpoint
-        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}')
+            elif distributed_backend is None:
+                m = 'When using multiple GPUs set ' \
+                    'Trainer(distributed_backend=dp) (or ddp)'
+                raise MisconfigurationException(m)
+
+        # use ddp automatically if nb_gpu_nodes > 1
+        if nb_gpu_nodes > 1 and self.use_dp:  # pragma: no cover
+            self.use_ddp = True
+            self.use_dp = False
+            w = 'DataParallel does not support nb_gpu_nodes > 1. ' \
+                'Switching to DistributedDataParallel for you. ' \
+                'To silence this warning set distributed_backend=ddp'
+            warnings.warn(w)
+
+        print('gpu available: {}, used: {}'.format(torch.cuda.is_available(), self.on_gpu))
+
+    def __configure_slurm_ddp(self, nb_gpu_nodes):
+        self.is_slurm_managing_tasks = False
+
+        # extract SLURM flag vars
+        # whenever we have the correct number of tasks, we let slurm manage processes
+        # otherwise we launch the required number of processes
+        if self.use_ddp:
+            self.nb_requested_gpus = self.num_gpus * nb_gpu_nodes
+            self.nb_slurm_tasks = 0
+            try:
+                self.nb_slurm_tasks = int(os.environ['SLURM_NTASKS'])
+                self.is_slurm_managing_tasks = self.nb_slurm_tasks == self.nb_requested_gpus
+
+                # in interactive mode we don't manage tasks
+                job_name = os.environ['SLURM_JOB_NAME']
+                if job_name == 'bash':
+                    self.is_slurm_managing_tasks = False
+
+            except Exception:
+                # likely not on slurm, so set the slurm managed flag to false
+                self.is_slurm_managing_tasks = False
+
+    def __set_nvidia_flags(self, is_slurm_managing_tasks, data_parallel_device_ids):
+        if data_parallel_device_ids is None:
+            return
+
+        # set the correct cuda visible devices (using pci order)
+        os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
+
+        # when slurm is managing the task it sets the visible devices
+        if not is_slurm_managing_tasks:
+            if type(data_parallel_device_ids) is int:
+                id_str = ','.join(str(x) for x in list(range(data_parallel_device_ids)))
+                os.environ["CUDA_VISIBLE_DEVICES"] = id_str
+            else:
+                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"]}')
 
     @property
     def data_parallel(self):
@@ -324,7 +398,7 @@ class Trainer(TrainerIO):
 
     def __is_overriden(self, f_name):
         model = self.__get_model()
-        super_object = super(model.__class__, model)
+        super_object = LightningModule
 
         # when code pointers are different, it was overriden
         is_overriden = getattr(model, f_name).__code__ is not getattr(super_object, f_name).__code__
@@ -344,7 +418,7 @@ class Trainer(TrainerIO):
         tqdm_dic.update(self.tqdm_metrics)
 
         if self.on_gpu:
-            tqdm_dic['gpu'] = '{}'.format(self.current_gpu_name)
+            tqdm_dic['gpu'] = '{}'.format(torch.cuda.current_device())
 
         return tqdm_dic
 
@@ -370,8 +444,10 @@ class Trainer(TrainerIO):
             self.nb_val_batches = max(1, self.nb_val_batches)
 
         # determine number of test batches
-        self.nb_test_batches = len(self.test_dataloader) if self.test_dataloader is not None else 0
-        self.nb_test_batches = int(self.nb_test_batches * self.test_percent_check)
+        if self.test_dataloader is not None:
+            self.nb_test_batches = sum(len(dataloader) for dataloader in self.test_dataloader)
+            self.nb_test_batches = int(self.nb_test_batches * self.test_percent_check)
+            self.nb_test_batches = max(1, self.nb_test_batches)
 
         # determine when to check validation
         self.val_check_batch = int(self.nb_tng_batches * self.val_check_interval)
@@ -384,40 +460,47 @@ class Trainer(TrainerIO):
 
             self.tqdm_metrics[k] = v
 
-    def __validation_forward(self, model, data_batch, batch_i, dataloader_i):
+    def __evaluation_forward(self, model, data_batch, batch_i, dataloader_i, test=False):
         # make dataloader_i arg in validation_step optional
         args = [data_batch, batch_i]
-        if len(self.val_dataloader) > 1:
+
+        if test and len(self.test_dataloader) > 1:
             args.append(dataloader_i)
 
-        if self.use_ddp:
+        elif not test and len(self.val_dataloader) > 1:
+            args.append(dataloader_i)
+
+        # handle DP, DDP forward
+        if self.use_ddp or self.use_dp:
             output = model(*args)
-        elif self.use_dp:
-            output = model(*args)
-        elif self.single_gpu:
-            # put inputs on gpu manually
-            gpu_id = self.data_parallel_device_ids[0]
-            data_batch = self.transfer_batch_to_gpu(data_batch, gpu_id)
+            return output
+
+        # CPU, single GPU
+        if self.single_gpu:
+            # for single GPU put inputs on gpu manually
+            root_gpu = 0
+            if type(self.data_parallel_device_ids) is list:
+                root_gpu = self.data_parallel_device_ids[0]
+            data_batch = self.transfer_batch_to_gpu(data_batch, root_gpu)
             args[0] = data_batch
 
-            # do non dp, ddp step
-            output = model.validation_step(*args)
-
+        if test:
+            output = model.test_step(*args)
         else:
-            # CPU
             output = model.validation_step(*args)
 
         return output
 
-    def validate(self, model, dataloader, max_batches, dataloader_i):
+    def evaluate(self, model, dataloaders, max_batches, test=False):
         """
-        Run validation code
+        Run evaluation code
         :param model: PT model
-        :param dataloader: PT dataloader
+        :param dataloaders: list of PT dataloaders
         :param max_batches: Scalar
+        :param dataloader_i:
+        :param test: boolean
         :return:
         """
-
         # enable eval mode
         model.zero_grad()
         model.eval()
@@ -429,34 +512,41 @@ class Trainer(TrainerIO):
         outputs = []
 
         # run training
-        for batch_i, data_batch in enumerate(dataloader):
+        for dataloader_i, dl in enumerate(dataloaders):
+            dl_outputs = []
+            for batch_i, data_batch in enumerate(dl):
 
-            if data_batch is None:  # pragma: no cover
-                continue
+                if data_batch is None:  # pragma: no cover
+                    continue
 
-            # stop short when on fast_dev_run (sets max_batch=1)
-            if batch_i >= max_batches:
-                break
+                # stop short when on fast_dev_run (sets max_batch=1)
+                if batch_i >= max_batches:
+                    break
 
-            # -----------------
-            # RUN VALIDATION STEP
-            # -----------------
-            output = self.__validation_forward(model, data_batch, batch_i, dataloader_i)
+                # -----------------
+                # RUN EVALUATION STEP
+                # -----------------
+                output = self.__evaluation_forward(model, data_batch, batch_i, dataloader_i,
+                                                   test)
 
-            # track outputs for collation
-            outputs.append(output)
+                # track outputs for collation
+                dl_outputs.append(output)
 
-            # batch done
-            if self.show_progress_bar:
-                self.progress_bar.update(1)
+                # batch done
+                if self.show_progress_bar:
+                    self.progress_bar.update(1)
+            outputs.append(dl_outputs)
+
+        eval_results = {}
 
         # give model a chance to do something with the outputs (and method defined)
-        val_results = {}
-        if self.__is_overriden('validation_end'):
-            if self.data_parallel:
-                val_results = model.module.validation_end(outputs)
-            else:
-                val_results = model.validation_end(outputs)
+        model = self.__get_model()
+        if len(dataloaders) == 1:
+            outputs = outputs[0]
+        if test and self.__is_overriden('test_end'):
+            eval_results = model.test_end(outputs)
+        elif self.__is_overriden('validation_end'):
+            eval_results = model.validation_end(outputs)
 
         # enable train mode again
         model.train()
@@ -464,7 +554,7 @@ class Trainer(TrainerIO):
         # enable gradients to save memory
         torch.set_grad_enabled(True)
 
-        return val_results
+        return eval_results
 
     def get_dataloaders(self, model):
         """
@@ -478,6 +568,10 @@ class Trainer(TrainerIO):
         self.val_dataloader = model.val_dataloader
 
         # handle returning an actual dataloader instead of a list of loaders
+        have_test_loaders = self.test_dataloader is not None
+        if have_test_loaders and not isinstance(self.test_dataloader, list):
+            self.test_dataloader = [self.test_dataloader]
+
         have_val_loaders = self.val_dataloader is not None
         if have_val_loaders and not isinstance(self.val_dataloader, list):
             self.val_dataloader = [self.val_dataloader]
@@ -503,9 +597,32 @@ class Trainer(TrainerIO):
 
         if self.use_ddp and self.val_dataloader is not None:
             for dataloader in self.val_dataloader:
-                if not isinstance(dataloader, DistributedSampler):
+                if not isinstance(dataloader.sampler, DistributedSampler):
                     msg = """
-                    Your val_dataloader(s) are not all DistributedSamplers.
+                    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.
+
+                    ie: this:
+                    dataset = myDataset()
+                    dataloader = Dataloader(dataset)
+
+                    becomes:
+                    dataset = myDataset()
+                    dist_sampler = torch.utils.data.distributed.DistributedSampler(dataset)
+                    dataloader = Dataloader(dataset, sampler=dist_sampler)
+
+                    If you want each process to load the full dataset, ignore this warning.
+                    """
+                    warnings.warn(msg)
+                    break
+
+        if self.use_ddp and self.test_dataloader is not None:
+            for dataloader in self.test_dataloader:
+                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.
@@ -528,7 +645,6 @@ class Trainer(TrainerIO):
     # MODEL TRAINING
     # -----------------------------
     def fit(self, model):
-
         # when using multi-node or DDP within a node start each module in a separate process
         if self.use_ddp:
             # must copy only the meta of the exp so it survives pickle/unpickle
@@ -550,7 +666,7 @@ class Trainer(TrainerIO):
                 If you're not using SLURM, ignore this message!
                 """
                 warnings.warn(msg)
-                mp.spawn(self.ddp_train, nprocs=len(self.data_parallel_device_ids), args=(model, ))
+                mp.spawn(self.ddp_train, nprocs=self.num_gpus, args=(model, ))
 
         # 1 gpu or dp option triggers training using DP module
         # easier to avoid NCCL issues
@@ -597,7 +713,7 @@ class Trainer(TrainerIO):
         # allow for lr schedulers as well
         self.optimizers, self.lr_schedulers = self.init_optimizers(model.configure_optimizers())
 
-        model.cuda(self.data_parallel_device_ids[0])
+        model.cuda(self.root_gpu)
 
         if self.use_amp:
             # An example
@@ -614,7 +730,7 @@ class Trainer(TrainerIO):
         # allow for lr schedulers as well
         self.optimizers, self.lr_schedulers = self.init_optimizers(model.configure_optimizers())
 
-        model.cuda(self.data_parallel_device_ids[0])
+        model.cuda(self.root_gpu)
 
         # check for this bug (amp + dp + !01 doesn't work)
         # https://github.com/NVIDIA/apex/issues/227
@@ -626,7 +742,12 @@ class Trainer(TrainerIO):
             """
             raise MisconfigurationException(m)
 
-        model = LightningDataParallel(model, device_ids=self.data_parallel_device_ids)
+        # create list of device ids
+        device_ids = self.data_parallel_device_ids
+        if type(device_ids) is int:
+            device_ids = list(range(device_ids))
+
+        model = LightningDataParallel(model, device_ids=device_ids)
 
         self.__run_pretrain_routine(model)
 
@@ -656,8 +777,8 @@ class Trainer(TrainerIO):
         self.show_progress_bar = self.show_progress_bar and self.node_rank == 0 and gpu_nb == 0
 
         # determine which process we are and world size
-        self.proc_rank = self.node_rank * len(self.data_parallel_device_ids) + gpu_nb
-        self.world_size = self.nb_gpu_nodes * len(self.data_parallel_device_ids)
+        self.proc_rank = self.node_rank * self.num_gpus + gpu_nb
+        self.world_size = self.nb_gpu_nodes * self.num_gpus
 
         # let the exp know the rank to avoid overwriting logs
         if self.experiment is not None:
@@ -677,6 +798,9 @@ class Trainer(TrainerIO):
         torch.cuda.set_device(gpu_nb)
         model.cuda(gpu_nb)
 
+        # override root GPU
+        self.root_gpu = gpu_nb
+
         # AMP
         # run through amp wrapper before going to distributed DP
         if self.use_amp:
@@ -740,6 +864,7 @@ class Trainer(TrainerIO):
         if self.data_parallel:
             ref_model = model.module
 
+        # give model convenience properties
         ref_model.trainer = self
 
         # set local properties on the model
@@ -747,6 +872,10 @@ class Trainer(TrainerIO):
         ref_model.use_dp = self.use_dp
         ref_model.use_ddp = self.use_ddp
         ref_model.use_amp = self.use_amp
+        ref_model.testing = self.testing
+
+        # register auto-resubmit when on SLURM
+        self.register_slurm_signal_handlers()
 
         # transfer data loaders from model
         self.get_dataloaders(ref_model)
@@ -758,43 +887,39 @@ class Trainer(TrainerIO):
         if self.proc_rank == 0 and self.print_weights_summary:
             ref_model.summarize()
 
-        # give model convenience properties
-        ref_model.trainer = self
-
+        # link up experiment object
         if self.experiment is not None:
             ref_model.experiment = self.experiment
 
-        # save exp to get started
-        if self.proc_rank == 0 and self.experiment is not None:
-            self.experiment.save()
+            # save exp to get started
+            if self.proc_rank == 0:
+                self.experiment.save()
 
         # track model now.
         # if cluster resets state, the model will update with the saved weights
         self.model = model
 
         # restore training and model before hpc call
-        self.restore_state_if_existing_checkpoint()
-
-        # enable cluster checkpointing
-        # also restores training state
-        # hpc checkpoint overrides any other checkpoints loaded before
-        if self.cluster is not None:  # pragma: no cover
-            self.enable_auto_hpc_walltime_manager()
+        self.restore_weights(model)
 
         # progress bar init
         if self.show_progress_bar:
             self.progress_bar = tqdm.tqdm(0, position=self.process_position)
 
-        # run tiny validation (if validation defined) to make sure program won't crash during val
+        # when testing requested only run test and return
+        if self.testing:
+            self.__run_evaluation(test=True)
+            return
+
+        # run tiny validation (if validation defined)
+        # to make sure program won't crash during val
         ref_model.on_sanity_check_start()
         if self.val_dataloader is not None and self.nb_sanity_val_steps > 0:
-            for ds_i, dataloader in enumerate(self.val_dataloader):
+            # reset progress_bar limit for sanity check
+            if self.show_progress_bar:
+                self.progress_bar.reset(self.nb_sanity_val_steps)
 
-                # reset progress_bar limit for sanity check
-                if self.show_progress_bar:
-                    self.progress_bar.reset(self.nb_sanity_val_steps)
-
-                self.validate(model, dataloader, self.nb_sanity_val_steps, ds_i)
+            self.evaluate(model, self.val_dataloader, self.nb_sanity_val_steps, self.testing)
 
         # ---------------------------
         # CORE TRAINING LOOP
@@ -804,6 +929,10 @@ class Trainer(TrainerIO):
     def __train(self):
         # run all epochs
         for epoch_nb in range(self.current_epoch, self.max_nb_epochs):
+            # set seed for distributed sampler (enables shuffling for each epoch)
+            if self.use_ddp:
+                self.tng_dataloader.sampler.set_epoch(epoch_nb)
+
             # get model
             model = self.__get_model()
 
@@ -871,8 +1000,10 @@ class Trainer(TrainerIO):
             # RUN VAL STEP
             # ---------------
             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:
-                self.__run_validation()
+                if can_check_epoch:
+                    self.__run_evaluation(test=self.testing)
 
             # when batch should be saved
             if (batch_nb + 1) % self.log_save_interval == 0 or early_stop_epoch:
@@ -888,7 +1019,7 @@ class Trainer(TrainerIO):
                 metrics = self.__tng_tqdm_dic
 
                 # add gpu memory
-                if self.on_gpu:
+                if self.on_gpu and self.log_gpu_memory:
                     mem_map = get_gpu_memory_map()
                     metrics.update(mem_map)
 
@@ -917,6 +1048,13 @@ class Trainer(TrainerIO):
             model = self.__get_model()
             model.on_epoch_end()
 
+    def test(self, model=None):
+        if model is not None:
+            self.testing = True
+            self.fit(model)
+        else:
+            self.__run_evaluation(test=True)
+
     def __metrics_to_scalars(self, metrics, blacklist=set()):
         new_metrics = {}
         for k, v in metrics.items():
@@ -937,10 +1075,13 @@ class Trainer(TrainerIO):
         return blacklist
 
     def transfer_batch_to_gpu(self, batch, gpu_id):
-        # base case
-        if isinstance(batch, torch.Tensor):
+        # base case: object can be directly moved using `cuda` or `to`
+        if callable(getattr(batch, 'cuda', None)):
             return batch.cuda(gpu_id)
 
+        elif callable(getattr(batch, 'to', None)):
+            return batch.to(torch.device('cuda', gpu_id))
+
         # when list
         elif isinstance(batch, list):
             for i, x in enumerate(batch):
@@ -984,7 +1125,9 @@ class Trainer(TrainerIO):
         elif self.use_dp:
             output = self.model(*args)
         elif self.single_gpu:
-            gpu_id = self.data_parallel_device_ids[0]
+            gpu_id = 0
+            if type(self.data_parallel_device_ids) is list:
+                gpu_id = self.data_parallel_device_ids[0]
             data_batch = self.transfer_batch_to_gpu(data_batch, gpu_id)
             args[0] = data_batch
             output = self.model.training_step(*args)
@@ -1000,7 +1143,7 @@ class Trainer(TrainerIO):
 
             # reduce prog metrics for tqdm when using dp
             if self.use_dp:
-                nb_gpus = len(self.data_parallel_device_ids)
+                nb_gpus = self.num_gpus
                 prog_output = reduce_distributed_output(prog_output, nb_gpus)
 
             model_specific_tqdm_metrics_dic = prog_output
@@ -1020,7 +1163,7 @@ class Trainer(TrainerIO):
 
         # when using dp need to reduce the loss
         if self.use_dp:
-            loss = reduce_distributed_output(loss, len(self.data_parallel_device_ids))
+            loss = reduce_distributed_output(loss, self.num_gpus)
 
         return loss, model_specific_tqdm_metrics_dic
 
@@ -1030,10 +1173,10 @@ class Trainer(TrainerIO):
             torch.nn.utils.clip_grad_norm_(model.parameters(), self.gradient_clip)
 
     def __print_nan_grads(self):
-        if self.print_nan_grads:
-            model = self.__get_model()
-            for param in model.parameters():
-                print(param.grad.float().sum())
+        model = self.__get_model()
+        for param in model.parameters():
+            if torch.isnan(param.grad.float()).any():
+                print(param, param.grad)
 
     def __run_tng_batch(self, data_batch, batch_nb):
         if data_batch is None:
@@ -1076,7 +1219,8 @@ class Trainer(TrainerIO):
                 model_ref.on_after_backward()
 
             # nan grads
-            self.__print_nan_grads()
+            if self.print_nan_grads:
+                self.__print_nan_grads()
 
             # track total loss for logging (avoid mem leaks)
             self.batch_loss_value += loss.item()
@@ -1109,33 +1253,48 @@ class Trainer(TrainerIO):
 
         return 0
 
-    def __run_validation(self):
-        # decide if can check epochs
-        can_check_epoch = (self.current_epoch + 1) % self.check_val_every_n_epoch == 0
-        if self.fast_dev_run:
-            print('skipping to check performance bc of --fast_dev_run')
-        elif not can_check_epoch:
-            return
+    def __run_evaluation(self, test=False):
+        # when testing make sure user defined a test step
+        can_run_test_step = False
+        if test:
+            can_run_test_step = self.__is_overriden('test_step') and self.__is_overriden('test_end')
+            if not can_run_test_step:
+                m = '''You called .test() without defining a test step or test_end.
+                Please define and try again'''
+                raise MisconfigurationException(m)
 
         # validate only if model has validation_step defined
-        if self.__is_overriden('validation_step'):
+        # test only if test_step or validation_step are defined
+        run_val_step = self.__is_overriden('validation_step')
+
+        if run_val_step or can_run_test_step:
 
             # hook
-            if self.__is_function_implemented('on_pre_performance_check'):
-                model = self.__get_model()
-                model.on_pre_performance_check()
+            model = self.__get_model()
+            model.on_pre_performance_check()
 
-            # use val_percent_check set on end of epoch
-            # use a small portion otherwise
-            max_batches = self.nb_val_batches if not self.fast_dev_run else 1
-            for ds_i, dataloader in enumerate(self.val_dataloader):
-                val_out_metrics = self.validate(self.model, dataloader, max_batches, ds_i)
-                self.__add_tqdm_metrics(val_out_metrics)
+            # select dataloaders
+            dataloaders = self.val_dataloader
+            max_batches = self.nb_val_batches
+
+            # calculate max batches to use
+            if test:
+                dataloaders = self.test_dataloader
+                max_batches = self.nb_test_batches
+
+            # cap max batches to 1 when using fast_dev_run
+            if self.fast_dev_run:
+                max_batches = 1
+
+            eval_out_metrics = self.evaluate(self.model,
+                                             dataloaders,
+                                             max_batches,
+                                             test)
+
+            self.__add_tqdm_metrics(eval_out_metrics)
 
             # hook
-            if self.__is_function_implemented('on_post_performance_check'):
-                model = self.__get_model()
-                model.on_post_performance_check()
+            model.on_post_performance_check()
 
             if self.show_progress_bar:
                 # add model specific metrics
@@ -1143,7 +1302,7 @@ class Trainer(TrainerIO):
                 self.progress_bar.set_postfix(**tqdm_metrics)
 
         # model checkpointing
-        if self.proc_rank == 0 and self.checkpoint_callback is not None:
+        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.__tng_tqdm_dic)

pytorch_lightning/trainer/trainer_io.py

@@ -0,0 +1,325 @@
+import os
+import re
+import signal
+import pdb
+from subprocess import call
+
+import torch
+
+from pytorch_lightning.pt_overrides.override_data_parallel import (
+    LightningDistributedDataParallel, LightningDataParallel)
+
+
+class TrainerIO(object):
+
+    def __get_model(self):
+        is_dp_module = isinstance(self.model, (LightningDistributedDataParallel,
+                                               LightningDataParallel))
+        model = self.model.module if is_dp_module else self.model
+        return model
+
+    # --------------------
+    # CHECK-POINTING
+    # --------------------
+    def restore_weights(self, model):
+        """
+        To restore weights we have two cases.
+        First, if we use the same experiment version, then restore the latest ckpt.
+        AFTER that, if we find weights from hpc checkpoint, then restore that.
+        :param model:
+        :return:
+        """
+        # restore weights if same exp version
+        self.restore_state_if_checkpoint_exists(model)
+
+        # if script called from hpc resubmit, load weights
+        self.restore_hpc_weights_if_needed(model)
+
+    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
+        if no_ckpt_callback or not os.path.exists(self.checkpoint_callback.filepath):
+            return
+
+        # restore trainer state and model if there is a weight for this experiment
+        last_epoch = -1
+        last_ckpt_name = None
+
+        # find last epoch
+        checkpoints = os.listdir(self.checkpoint_callback.filepath)
+        for name in checkpoints:
+            # ignore hpc ckpts
+            if 'hpc_' in name:
+                continue
+
+            if '.ckpt' in name:
+                epoch = name.split('epoch_')[1]
+                epoch = int(re.sub('[^0-9]', '', epoch))
+
+                if epoch > last_epoch:
+                    last_epoch = epoch
+                    last_ckpt_name = name
+
+        # restore last checkpoint
+        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}')
+
+    # --------------------
+    # HPC SIGNAL HANDLING
+    # --------------------
+    def register_slurm_signal_handlers(self):
+        # see if we're using slurm (not interactive)
+        on_slurm = False
+        try:
+            job_name = os.environ['SLURM_JOB_NAME']
+            if job_name != 'bash':
+                on_slurm = True
+        except Exception as e:
+            pass
+
+        if on_slurm:
+            print('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')
+            self.hpc_save(self.weights_save_path, self.experiment)
+
+            # find job id
+            job_id = os.environ['SLURM_JOB_ID']
+            cmd = 'scontrol requeue {}'.format(job_id)
+
+            # requeue job
+            print('\nrequeing job {}...'.format(job_id))
+            result = call(cmd, shell=True)
+
+            # print result text
+            if result == 0:
+                print('requeued exp ', job_id)
+            else:
+                print('requeue failed...')
+
+            # close experiment to avoid issues
+            self.experiment.close()
+
+    def term_handler(self, signum, frame):
+        # save
+        print("bypassing sigterm")
+
+    # --------------------
+    # MODEL SAVE CHECKPOINT
+    # --------------------
+    def save_checkpoint(self, filepath):
+        checkpoint = self.dump_checkpoint()
+
+        # do the actual save
+        torch.save(checkpoint, filepath)
+
+    def restore(self, checkpoint_path, on_gpu):
+
+        # if on_gpu:
+        #     checkpoint = torch.load(checkpoint_path)
+        # else:
+        # load on CPU first
+        checkpoint = torch.load(checkpoint_path, map_location=lambda storage, loc: storage)
+
+        # load model state
+        model = self.__get_model()
+
+        # load the state_dict on the model automatically
+        model.load_state_dict(checkpoint['state_dict'])
+        if on_gpu:
+            model.cuda(self.root_gpu)
+
+        # load training state (affects trainer only)
+        self.restore_training_state(checkpoint)
+
+    def dump_checkpoint(self):
+
+        checkpoint = {
+            'epoch': self.current_epoch,
+            'global_step': self.global_step
+        }
+
+        if self.checkpoint_callback is not None:
+            checkpoint['checkpoint_callback_best'] = self.checkpoint_callback.best
+
+        if self.early_stop_callback is not None:
+            checkpoint['early_stop_callback_wait'] = self.early_stop_callback.wait
+            checkpoint['early_stop_callback_patience'] = self.early_stop_callback.patience
+
+        # save optimizers
+        optimizer_states = []
+        for i, optimizer in enumerate(self.optimizers):
+            optimizer_states.append(optimizer.state_dict())
+
+        checkpoint['optimizer_states'] = optimizer_states
+
+        # save lr schedulers
+        lr_schedulers = []
+        for i, scheduler in enumerate(self.lr_schedulers):
+            lr_schedulers.append(scheduler.state_dict())
+
+        checkpoint['lr_schedulers'] = lr_schedulers
+
+        # add the state_dict from the model
+        model = self.__get_model()
+        checkpoint['state_dict'] = model.state_dict()
+
+        # give the model a chance to add a few things
+        model.on_save_checkpoint(checkpoint)
+
+        return checkpoint
+
+    # --------------------
+    # HPC IO
+    # --------------------
+    def restore_hpc_weights_if_needed(self, model):
+        """
+        If there is a set of hpc weights, use as signal to restore model
+        :param model:
+        :return:
+        """
+        # look for hpc weights
+        folderpath = self.weights_save_path
+        if os.path.exists(folderpath):
+            files = os.listdir(folderpath)
+            hpc_weight_paths = [x for x in files if 'hpc_ckpt' in x]
+
+            # if hpc weights exist restore model
+            if len(hpc_weight_paths) > 0:
+                self.hpc_load(folderpath, self.on_gpu)
+
+    def restore_training_state(self, checkpoint):
+        """
+        Restore trainer state.
+        Model will get its change to update
+        :param checkpoint:
+        :return:
+        """
+        if self.checkpoint_callback is not None:
+            self.checkpoint_callback.best = checkpoint['checkpoint_callback_best']
+
+        if self.early_stop_callback is not None:
+            self.early_stop_callback.wait = checkpoint['early_stop_callback_wait']
+            self.early_stop_callback.patience = checkpoint['early_stop_callback_patience']
+
+        self.global_step = checkpoint['global_step']
+        self.current_epoch = checkpoint['epoch']
+
+        # restore the optimizers
+        optimizer_states = checkpoint['optimizer_states']
+        for optimizer, opt_state in zip(self.optimizers, optimizer_states):
+            optimizer.load_state_dict(opt_state)
+
+            # move optimizer to GPU 1 weight at a time
+            # avoids OOM
+            if self.root_gpu is not None:
+                for state in optimizer.state.values():
+                    for k, v in state.items():
+                        if isinstance(v, torch.Tensor):
+                            state[k] = v.cuda(self.root_gpu)
+
+        # restore the lr schedulers
+        lr_schedulers = checkpoint['lr_schedulers']
+        for scheduler, lrs_state in zip(self.lr_schedulers, lr_schedulers):
+            scheduler.load_state_dict(lrs_state)
+
+    # ----------------------------------
+    # PRIVATE OPS
+    # ----------------------------------
+    def hpc_save(self, folderpath, experiment):
+        # make sure the checkpoint folder exists
+        os.makedirs(folderpath, exist_ok=True)
+
+        # save exp to make sure we get all the metrics
+        experiment.save()
+
+        ckpt_number = self.max_ckpt_in_folder(folderpath) + 1
+
+        if not os.path.exists(folderpath):
+            os.makedirs(folderpath, exist_ok=True)
+        filepath = '{}/hpc_ckpt_{}.ckpt'.format(folderpath, ckpt_number)
+
+        # give model a chance to do something on hpc_save
+        model = self.__get_model()
+        checkpoint = self.dump_checkpoint()
+
+        model.on_hpc_save(checkpoint)
+
+        # do the actual save
+        torch.save(checkpoint, filepath)
+
+        return filepath
+
+    def hpc_load(self, folderpath, on_gpu):
+        filepath = '{}/hpc_ckpt_{}.ckpt'.format(folderpath, self.max_ckpt_in_folder(folderpath))
+
+        # load on CPU first
+        checkpoint = torch.load(filepath, map_location=lambda storage, loc: storage)
+
+        # load model state
+        model = self.__get_model()
+
+        # load the state_dict on the model automatically
+        model.load_state_dict(checkpoint['state_dict'])
+
+        if self.root_gpu is not None:
+            model.cuda(self.root_gpu)
+
+        # load training state (affects trainer only)
+        self.restore_training_state(checkpoint)
+
+        # call model hook
+        model.on_hpc_load(checkpoint)
+
+        print(f'restored hpc model from: {filepath}')
+
+    def max_ckpt_in_folder(self, path, name_key='ckpt_'):
+        files = os.listdir(path)
+        files = [x for x in files if name_key in x]
+        if len(files) == 0:
+            return 0
+
+        ckpt_vs = []
+        for name in files:
+            name = name.split(name_key)[-1]
+            name = re.sub('[^0-9]', '', name)
+            ckpt_vs.append(int(name))
+
+        return max(ckpt_vs)
+
+
+def load_hparams_from_tags_csv(tags_csv):
+    from argparse import Namespace
+    import pandas as pd
+
+    tags_df = pd.read_csv(tags_csv)
+    dic = tags_df.to_dict(orient='records')
+
+    ns_dict = {row['key']: convert(row['value']) for row in dic}
+
+    ns = Namespace(**ns_dict)
+    return ns
+
+
+def convert(val):
+    constructors = [int, float, str]
+
+    if type(val) is str:
+        if val.lower() == 'true':
+            return True
+        if val.lower() == 'false':
+            return False
+
+    for c in constructors:
+        try:
+            return c(val)
+        except ValueError:
+            pass
+    return val

setup.py

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

tests/debug.py

@@ -1,5 +1,6 @@
 from pytorch_lightning import Trainer
 from examples import LightningTemplateModel
+from pytorch_lightning.testing import LightningTestModel
 from argparse import Namespace
 from test_tube import Experiment
 from pytorch_lightning.callbacks import ModelCheckpoint
@@ -12,6 +13,7 @@ from torch.nn import functional as F
 from torch.utils.data import DataLoader
 from torchvision.datasets import MNIST
 import numpy as np
+import pdb
 
 
 class CoolModel(pl.LightningModule):
@@ -73,10 +75,11 @@ def get_model():
     return model, hparams
 
 
-def get_exp(debug=True):
+def get_exp(debug=True, version=None):
     # set up exp object without actually saving logs
     root_dir = os.path.dirname(os.path.realpath(__file__))
-    exp = Experiment(debug=debug, save_dir=root_dir, name='tests_tt_dir')
+    save_dir = os.path.join(root_dir, 'save_dir')
+    exp = Experiment(debug=debug, save_dir=save_dir, name='tests_tt_dir', version=version)
     return exp
 
 
@@ -99,7 +102,7 @@ def clear_save_dir():
         shutil.rmtree(save_dir)
 
 
-def load_model(exp, save_dir):
+def load_model(exp, save_dir, on_gpu, map_location=None, module_class=LightningTemplateModel):
 
     # load trained model
     tags_path = exp.get_data_path(exp.name, exp.version)
@@ -108,8 +111,10 @@ def load_model(exp, save_dir):
     checkpoints = [x for x in os.listdir(save_dir) if '.ckpt' in x]
     weights_dir = os.path.join(save_dir, checkpoints[0])
 
-    trained_model = LightningTemplateModel.load_from_metrics(weights_path=weights_dir,
-                                                             tags_csv=tags_path, on_gpu=True)
+    trained_model = module_class.load_from_metrics(weights_path=weights_dir,
+                                                   tags_csv=tags_path,
+                                                   on_gpu=on_gpu,
+                                                   map_location=map_location)
 
     assert trained_model is not None, 'loading model failed'
 
@@ -131,12 +136,10 @@ def run_prediction(dataloader, trained_model):
     val_acc = torch.sum(y == labels_hat).item() / (len(y) * 1.0)
     val_acc = torch.tensor(val_acc)
     val_acc = val_acc.item()
-
-    print(val_acc)
-
     assert val_acc > 0.70, 'this model is expected to get > 0.7 in test set (it got %f)' % val_acc
 
 
+# ------------------------------------------------------------------------
 def run_gpu_model_test(trainer_options, model, hparams, on_gpu=True):
     save_dir = init_save_dir()
 
@@ -162,7 +165,7 @@ def run_gpu_model_test(trainer_options, model, hparams, on_gpu=True):
     # test model loading
     pretrained_model = load_model(exp, save_dir, on_gpu)
 
-    # test model preds
+    # test new model accuracy
     run_prediction(model.test_dataloader, pretrained_model)
 
     if trainer.use_ddp:
@@ -177,19 +180,117 @@ def run_gpu_model_test(trainer_options, model, hparams, on_gpu=True):
     clear_save_dir()
 
 
+def assert_ok_val_acc(trainer):
+    # this model should get 0.80+ acc
+    acc = trainer.tng_tqdm_dic['val_acc']
+    assert acc > 0.50, f'model failed to get expected 0.50 validation accuracy. Got: {acc}'
+
+
+def assert_ok_test_acc(trainer):
+    # this model should get 0.80+ acc
+    acc = trainer.tng_tqdm_dic['test_acc']
+    assert acc > 0.50, f'model failed to get expected 0.50 validation accuracy. Got: {acc}'
+
+
+def get_hparams(continue_training=False, hpc_exp_number=0):
+    root_dir = os.path.dirname(os.path.realpath(__file__))
+
+    args = {
+        'drop_prob': 0.2,
+        'batch_size': 32,
+        'in_features': 28 * 28,
+        'learning_rate': 0.001 * 8,
+        'optimizer_name': 'adam',
+        'data_root': os.path.join(root_dir, 'mnist'),
+        'out_features': 10,
+        'hidden_dim': 1000}
+
+    if continue_training:
+        args['test_tube_do_checkpoint_load'] = True
+        args['hpc_exp_number'] = hpc_exp_number
+
+    hparams = Namespace(**args)
+    return hparams
+
+
 def main():
+    """
+    Make sure DDP + AMP continue training correctly
+    :return:
+    """
+    hparams = get_hparams()
+    model = LightningTestModel(hparams)
 
-    os.environ['MASTER_PORT'] = str(np.random.randint(12000, 19000, 1)[0])
-    model, hparams = get_model()
     trainer_options = dict(
-        max_nb_epochs=1,
-        train_percent_check=0.4,
-        val_percent_check=0.2,
-        gpus=[0, 1],
-        distributed_backend='ddp'
+        show_progress_bar=True,
+        max_nb_epochs=4,
+        gpus=2,
+        distributed_backend='dp',
     )
 
-    run_gpu_model_test(trainer_options, model, hparams)
+    save_dir = init_save_dir()
+
+    # exp file to get meta
+    exp = get_exp(False)
+    exp.argparse(hparams)
+    exp.save()
+
+    # exp file to get weights
+    checkpoint = ModelCheckpoint(save_dir)
+
+    # add these to the trainer options
+    trainer_options['experiment'] = exp
+    trainer_options['checkpoint_callback'] = checkpoint
+
+    # fit model
+    trainer = Trainer(**trainer_options)
+    trainer.is_slurm_managing_tasks = True
+    result = trainer.fit(model)
+
+    # track epoch before saving
+    real_global_epoch = trainer.current_epoch
+
+    # correct result and ok accuracy
+    assert result == 1, 'amp + dp model failed to complete'
+
+    # ---------------------------
+    # HPC LOAD/SAVE
+    # ---------------------------
+    # save
+    trainer.hpc_save(save_dir, exp)
+
+    # init new trainer
+    new_exp = get_exp(False, version=exp.version)
+    trainer_options['experiment'] = new_exp
+    trainer_options['checkpoint_callback'] = ModelCheckpoint(save_dir)
+    trainer_options['train_percent_check'] = 0.2
+    trainer_options['val_percent_check'] = 0.2
+    trainer_options['max_nb_epochs'] = 1
+    new_trainer = Trainer(**trainer_options)
+
+    # set the epoch start hook so we can predict before the model does the full training
+    def assert_good_acc():
+        assert trainer.current_epoch == real_global_epoch and trainer.current_epoch > 0
+
+        # if model and state loaded correctly, predictions will be good even though we
+        # haven't trained with the new loaded model
+        dp_model = new_trainer.model
+        dp_model.eval()
+
+        _ = [run_prediction(dataloader, dp_model, dp=True) for dataloader in trainer.val_dataloader]
+
+    # new model
+    model = LightningTestModel(hparams)
+    model.on_sanity_check_start = assert_good_acc
+
+    # fit new model which should load hpc weights
+    new_trainer.fit(model)
+
+    # test freeze on gpu
+    model.freeze()
+    model.unfreeze()
+
+    clear_save_dir()
 
 
 if __name__ == '__main__':

tests/test_models.py

@@ -6,11 +6,19 @@ from argparse import Namespace
 import pytest
 import numpy as np
 import torch
-from test_tube import Experiment, SlurmCluster
+from test_tube import Experiment
 
 # sys.path += [os.path.abspath('..'), os.path.abspath('../..')]
 from pytorch_lightning import Trainer
-from pytorch_lightning.testing import LightningTestModel, NoValEndTestModel, NoValModel
+from pytorch_lightning.testing import (
+    LightningTestModel,
+    LightningTestModelBase,
+    LightningValidationMixin,
+    LightningValidationStepMixin,
+    LightningValidationMultipleDataloadersMixin,
+    LightningTestMixin,
+    LightningTestMultipleDataloadersMixin,
+)
 from pytorch_lightning.callbacks import (
     ModelCheckpoint,
     EarlyStopping,
@@ -18,8 +26,9 @@ from pytorch_lightning.callbacks import (
 )
 from pytorch_lightning.utilities.debugging import MisconfigurationException
 from pytorch_lightning.root_module import memory
-from pytorch_lightning.models.trainer import reduce_distributed_output
+from pytorch_lightning.trainer.trainer import reduce_distributed_output
 from pytorch_lightning.root_module import model_saving
+from pytorch_lightning.trainer import trainer_io
 from examples import LightningTemplateModel
 
 SEED = 2334
@@ -30,6 +39,303 @@ np.random.seed(SEED)
 # ------------------------------------------------------------------------
 # TESTS
 # ------------------------------------------------------------------------
+def test_dp_resume():
+    """
+    Make sure DP continues training correctly
+    :return:
+    """
+    if not can_run_gpu_test():
+        return
+
+    hparams = get_hparams()
+    model = LightningTestModel(hparams)
+
+    trainer_options = dict(
+        show_progress_bar=True,
+        max_nb_epochs=2,
+        gpus=2,
+        distributed_backend='dp',
+    )
+
+    save_dir = init_save_dir()
+
+    # exp file to get meta
+    exp = get_exp(False)
+    exp.argparse(hparams)
+    exp.save()
+
+    # exp file to get weights
+    checkpoint = ModelCheckpoint(save_dir)
+
+    # add these to the trainer options
+    trainer_options['experiment'] = exp
+    trainer_options['checkpoint_callback'] = checkpoint
+
+    # fit model
+    trainer = Trainer(**trainer_options)
+    trainer.is_slurm_managing_tasks = True
+    result = trainer.fit(model)
+
+    # track epoch before saving
+    real_global_epoch = trainer.current_epoch
+
+    # correct result and ok accuracy
+    assert result == 1, 'amp + dp model failed to complete'
+
+    # ---------------------------
+    # HPC LOAD/SAVE
+    # ---------------------------
+    # save
+    trainer.hpc_save(save_dir, exp)
+
+    # init new trainer
+    new_exp = get_exp(False, version=exp.version)
+    trainer_options['experiment'] = new_exp
+    trainer_options['checkpoint_callback'] = ModelCheckpoint(save_dir)
+    trainer_options['train_percent_check'] = 0.2
+    trainer_options['val_percent_check'] = 0.2
+    trainer_options['max_nb_epochs'] = 1
+    new_trainer = Trainer(**trainer_options)
+
+    # set the epoch start hook so we can predict before the model does the full training
+    def assert_good_acc():
+        assert new_trainer.current_epoch == real_global_epoch and new_trainer.current_epoch > 0
+
+        # if model and state loaded correctly, predictions will be good even though we
+        # haven't trained with the new loaded model
+        dp_model = new_trainer.model
+        dp_model.eval()
+
+        _ = [run_prediction(dataloader, dp_model, dp=True) for dataloader in trainer.val_dataloader]
+
+    # new model
+    model = LightningTestModel(hparams)
+    model.on_sanity_check_start = assert_good_acc
+
+    # fit new model which should load hpc weights
+    new_trainer.fit(model)
+
+    # test freeze on gpu
+    model.freeze()
+    model.unfreeze()
+
+    clear_save_dir()
+
+
+def test_running_test_pretrained_model_ddp():
+    """Verify test() on pretrained model"""
+    if not can_run_gpu_test():
+        return
+
+    hparams = get_hparams()
+    model = LightningTestModel(hparams)
+
+    save_dir = init_save_dir()
+
+    # exp file to get meta
+    exp = get_exp(False)
+    exp.argparse(hparams)
+    exp.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,
+        experiment=exp,
+        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(exp, save_dir, on_gpu=True, module_class=LightningTestModel)
+
+    # run test set
+    new_trainer = Trainer(**trainer_options)
+    new_trainer.test(pretrained_model)
+
+    run_prediction(model.test_dataloader, pretrained_model)
+
+    # test we have good test accuracy
+    clear_save_dir()
+
+
+def test_running_test_after_fitting():
+    """Verify test() on fitted model"""
+    hparams = get_hparams()
+    model = LightningTestModel(hparams)
+
+    save_dir = init_save_dir()
+
+    # exp file to get meta
+    exp = get_exp(False)
+    exp.argparse(hparams)
+    exp.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,
+        test_percent_check=0.2,
+        checkpoint_callback=checkpoint,
+        experiment=exp
+    )
+
+    # fit model
+    trainer = Trainer(**trainer_options)
+    result = trainer.fit(model)
+
+    assert result == 1, 'training failed to complete'
+
+    trainer.test()
+
+    # test we have good test accuracy
+    assert_ok_test_acc(trainer)
+
+    clear_save_dir()
+
+
+def test_running_test_without_val():
+    """Verify test() works on a model with no val_loader"""
+    class CurrentTestModel(LightningTestMixin, LightningTestModelBase):
+        pass
+    hparams = get_hparams()
+    model = CurrentTestModel(hparams)
+
+    save_dir = init_save_dir()
+
+    # exp file to get meta
+    exp = get_exp(False)
+    exp.argparse(hparams)
+    exp.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,
+        test_percent_check=0.2,
+        checkpoint_callback=checkpoint,
+        experiment=exp
+    )
+
+    # fit model
+    trainer = Trainer(**trainer_options)
+    result = trainer.fit(model)
+
+    assert result == 1, 'training failed to complete'
+
+    trainer.test()
+
+    # test we have good test accuracy
+    assert_ok_test_acc(trainer)
+
+    clear_save_dir()
+
+
+def test_running_test_pretrained_model():
+    """Verify test() on pretrained model"""
+    hparams = get_hparams()
+    model = LightningTestModel(hparams)
+
+    save_dir = init_save_dir()
+
+    # exp file to get meta
+    exp = get_exp(False)
+    exp.argparse(hparams)
+    exp.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,
+        experiment=exp
+    )
+
+    # 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(
+        exp, save_dir, on_gpu=False, module_class=LightningTestModel
+    )
+
+    new_trainer = Trainer(**trainer_options)
+    new_trainer.test(pretrained_model)
+
+    # test we have good test accuracy
+    assert_ok_test_acc(new_trainer)
+    clear_save_dir()
+
+
+def test_running_test_pretrained_model_dp():
+    """Verify test() on pretrained model"""
+    if not can_run_gpu_test():
+        return
+
+    hparams = get_hparams()
+    model = LightningTestModel(hparams)
+
+    save_dir = init_save_dir()
+
+    # exp file to get meta
+    exp = get_exp(False)
+    exp.argparse(hparams)
+    exp.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,
+        experiment=exp,
+        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 = load_model(exp, save_dir, on_gpu=True, module_class=LightningTestModel)
+
+    new_trainer = Trainer(**trainer_options)
+    new_trainer.test(pretrained_model)
+
+    # test we have good test accuracy
+    assert_ok_test_acc(new_trainer)
+    clear_save_dir()
+
+
 def test_gradient_accumulation_scheduling():
     """
     Test grad accumulation by the freq of optimizer updates
@@ -107,13 +413,7 @@ def test_multi_gpu_model_ddp():
     Make sure DDP works
     :return:
     """
-    if not torch.cuda.is_available():
-        warnings.warn('test_multi_gpu_model_ddp cannot run.'
-                      ' Rerun on a GPU node to run this test')
-        return
-    if not torch.cuda.device_count() > 1:
-        warnings.warn('test_multi_gpu_model_ddp cannot run.'
-                      ' Rerun on a node with 2+ GPUs to run this test')
+    if not can_run_gpu_test():
         return
 
     os.environ['MASTER_PORT'] = str(np.random.randint(12000, 19000, 1)[0])
@@ -161,13 +461,7 @@ def test_optimizer_return_options():
 
 
 def test_single_gpu_batch_parse():
-    if not torch.cuda.is_available():
-        warnings.warn('test_amp_gpu_ddp cannot run.'
-                      'Rerun on a GPU node to run this test')
-        return
-    if not torch.cuda.device_count() > 1:
-        warnings.warn('test_amp_gpu_ddp cannot run.'
-                      'Rerun on a node with 2+ GPUs to run this test')
+    if not can_run_gpu_test():
         return
 
     trainer = Trainer()
@@ -241,7 +535,10 @@ def test_no_val_module():
     :return:
     """
     hparams = get_hparams()
-    model = NoValModel(hparams)
+
+    class CurrentTestModel(LightningTestModelBase):
+        pass
+    model = CurrentTestModel(hparams)
 
     save_dir = init_save_dir()
 
@@ -252,7 +549,6 @@ def test_no_val_module():
 
     trainer_options = dict(
         max_nb_epochs=1,
-        cluster=SlurmCluster(),
         experiment=exp,
         checkpoint_callback=ModelCheckpoint(save_dir)
     )
@@ -261,7 +557,7 @@ def test_no_val_module():
     trainer = Trainer(**trainer_options)
     result = trainer.fit(model)
 
-    # traning complete
+    # training complete
     assert result == 1, 'amp + ddp model failed to complete'
 
     # save model
@@ -284,8 +580,11 @@ def test_no_val_end_module():
     Tests use case where trainer saves the model, and user loads it from tags independently
     :return:
     """
+
+    class CurrentTestModel(LightningValidationStepMixin, LightningTestModelBase):
+        pass
     hparams = get_hparams()
-    model = NoValEndTestModel(hparams)
+    model = CurrentTestModel(hparams)
 
     save_dir = init_save_dir()
 
@@ -296,7 +595,6 @@ def test_no_val_end_module():
 
     trainer_options = dict(
         max_nb_epochs=1,
-        cluster=SlurmCluster(),
         experiment=exp,
         checkpoint_callback=ModelCheckpoint(save_dir)
     )
@@ -370,14 +668,43 @@ def test_amp_single_gpu():
     trainer_options = dict(
         show_progress_bar=True,
         max_nb_epochs=1,
-        gpus=[0],
-        distributed_backend='dp',
+        gpus=1,
+        distributed_backend='ddp',
         use_amp=True
     )
 
     run_gpu_model_test(trainer_options, model, hparams)
 
 
+def test_no_amp_single_gpu():
+    """
+    Make sure DDP + AMP work
+    :return:
+    """
+    if not torch.cuda.is_available():
+        warnings.warn('test_amp_gpu_ddp cannot run.'
+                      'Rerun on a GPU node to run this test')
+        return
+    if not torch.cuda.device_count() > 1:
+        warnings.warn('test_amp_gpu_ddp cannot run.'
+                      'Rerun on a node with 2+ GPUs to run this test')
+        return
+
+    hparams = get_hparams()
+    model = LightningTestModel(hparams)
+
+    trainer_options = dict(
+        show_progress_bar=True,
+        max_nb_epochs=1,
+        gpus=1,
+        distributed_backend='dp',
+        use_amp=True
+    )
+
+    with pytest.raises((MisconfigurationException, ModuleNotFoundError)):
+        run_gpu_model_test(trainer_options, model, hparams)
+
+
 def test_cpu_restore_training():
     """
     Verify continue training session on CPU
@@ -449,13 +776,7 @@ def test_amp_gpu_ddp():
     Make sure DDP + AMP work
     :return:
     """
-    if not torch.cuda.is_available():
-        warnings.warn('test_amp_gpu_ddp cannot run.'
-                      'Rerun on a GPU node to run this test')
-        return
-    if not torch.cuda.device_count() > 1:
-        warnings.warn('test_amp_gpu_ddp cannot run.'
-                      'Rerun on a node with 2+ GPUs to run this test')
+    if not can_run_gpu_test():
         return
 
     os.environ['MASTER_PORT'] = str(np.random.randint(12000, 19000, 1)[0])
@@ -466,7 +787,7 @@ def test_amp_gpu_ddp():
     trainer_options = dict(
         show_progress_bar=True,
         max_nb_epochs=1,
-        gpus=[0, 1],
+        gpus=2,
         distributed_backend='ddp',
         use_amp=True
     )
@@ -489,10 +810,10 @@ def test_cpu_slurm_save_load():
     exp.argparse(hparams)
     exp.save()
 
-    cluster_a = SlurmCluster()
+    version = exp.version
+
     trainer_options = dict(
         max_nb_epochs=1,
-        cluster=cluster_a,
         experiment=exp,
         checkpoint_callback=ModelCheckpoint(save_dir)
     )
@@ -516,22 +837,18 @@ def test_cpu_slurm_save_load():
     model.eval()
     pred_before_saving = model(x)
 
-    # test registering a save function
-    trainer.enable_auto_hpc_walltime_manager()
-
     # test HPC saving
     # simulate snapshot on slurm
     saved_filepath = trainer.hpc_save(save_dir, exp)
     assert os.path.exists(saved_filepath)
 
-    # wipe-out trainer and model
-    # retrain with not much data... this simulates picking training back up after slurm
-    # we want to see if the weights come back correctly
-    continue_tng_hparams = get_hparams(continue_training=True,
-                                       hpc_exp_number=cluster_a.hpc_exp_number)
+    # new exp file to get meta
+    exp = get_exp(False, version=version)
+    exp.argparse(hparams)
+    exp.save()
+
     trainer_options = dict(
         max_nb_epochs=1,
-        cluster=SlurmCluster(continue_tng_hparams),
         experiment=exp,
         checkpoint_callback=ModelCheckpoint(save_dir),
     )
@@ -567,7 +884,7 @@ def test_loading_meta_tags():
 
     # load tags
     tags_path = exp.get_data_path(exp.name, exp.version) + '/meta_tags.csv'
-    tags = model_saving.load_hparams_from_tags_csv(tags_path)
+    tags = trainer_io.load_hparams_from_tags_csv(tags_path)
 
     assert tags.batch_size == 32 and tags.hidden_dim == 1000
 
@@ -612,7 +929,6 @@ def test_model_saving_loading():
 
     trainer_options = dict(
         max_nb_epochs=1,
-        cluster=SlurmCluster(),
         experiment=exp,
         checkpoint_callback=ModelCheckpoint(save_dir)
     )
@@ -667,13 +983,7 @@ def test_amp_gpu_ddp_slurm_managed():
     Make sure DDP + AMP work
     :return:
     """
-    if not torch.cuda.is_available():
-        warnings.warn('test_amp_gpu_ddp cannot run.'
-                      ' Rerun on a GPU node to run this test')
-        return
-    if not torch.cuda.device_count() > 1:
-        warnings.warn('test_amp_gpu_ddp cannot run.'
-                      ' Rerun on a node with 2+ GPUs to run this test')
+    if not can_run_gpu_test():
         return
 
     # simulate setting slurm flags
@@ -720,8 +1030,7 @@ def test_amp_gpu_ddp_slurm_managed():
     assert trainer.resolve_root_node_address('abc[23-24, 45-40, 40]') == 'abc23'
 
     # test model loading with a map_location
-    map_location = 'cuda:1'
-    pretrained_model = load_model(exp, save_dir, True, map_location)
+    pretrained_model = load_model(exp, save_dir, True)
 
     # test model preds
     run_prediction(model.test_dataloader, pretrained_model)
@@ -821,28 +1130,46 @@ def test_single_gpu_model():
         max_nb_epochs=1,
         train_percent_check=0.1,
         val_percent_check=0.1,
-        gpus=[0]
+        gpus=1
     )
 
     run_gpu_model_test(trainer_options, model, hparams)
 
 
+def test_multi_gpu_none_backend():
+    """
+    Make sure when using multiple GPUs the user can't use
+    distributed_backend = None
+    :return:
+    """
+    if not can_run_gpu_test():
+        return
+
+    model, hparams = get_model()
+    trainer_options = dict(
+        show_progress_bar=False,
+        max_nb_epochs=1,
+        train_percent_check=0.1,
+        val_percent_check=0.1,
+        gpus='-1'
+    )
+
+    with pytest.raises(MisconfigurationException):
+        run_gpu_model_test(trainer_options, model, hparams)
+
+
 def test_multi_gpu_model_dp():
     """
     Make sure DP works
     :return:
     """
-    if not torch.cuda.is_available():
-        warnings.warn('test_multi_gpu_model_dp cannot run.'
-                      ' Rerun on a GPU node to run this test')
-        return
-    if not torch.cuda.device_count() > 1:
-        warnings.warn('test_multi_gpu_model_dp cannot run.'
-                      ' Rerun on a node with 2+ GPUs to run this test')
+    if not can_run_gpu_test():
         return
+
     model, hparams = get_model()
     trainer_options = dict(
         show_progress_bar=False,
+        distributed_backend='dp',
         max_nb_epochs=1,
         train_percent_check=0.1,
         val_percent_check=0.1,
@@ -860,14 +1187,9 @@ def test_amp_gpu_dp():
     Make sure DP + AMP work
     :return:
     """
-    if not torch.cuda.is_available():
-        warnings.warn('test_amp_gpu_dp cannot run.'
-                      ' Rerun on a GPU node to run this test')
-        return
-    if not torch.cuda.device_count() > 1:
-        warnings.warn('test_amp_gpu_dp cannot run.'
-                      ' Rerun on a node with 2+ GPUs to run this test')
+    if not can_run_gpu_test():
         return
+
     model, hparams = get_model()
     trainer_options = dict(
         max_nb_epochs=1,
@@ -884,11 +1206,7 @@ def test_ddp_sampler_error():
     Make sure DDP + AMP work
     :return:
     """
-    if not torch.cuda.is_available():
-        warnings.warn('test_amp_gpu_ddp cannot run. Rerun on a GPU node to run this test')
-        return
-    if not torch.cuda.device_count() > 1:
-        warnings.warn('test_amp_gpu_ddp cannot run. Rerun on a node with 2+ GPUs to run this test')
+    if not can_run_gpu_test():
         return
 
     os.environ['MASTER_PORT'] = str(np.random.randint(12000, 19000, 1)[0])
@@ -919,16 +1237,19 @@ def test_multiple_val_dataloader():
     Verify multiple val_dataloader
     :return:
     """
+    class CurrentTestModel(
+        LightningValidationMultipleDataloadersMixin,
+        LightningTestModelBase
+    ):
+        pass
     hparams = get_hparams()
-    model = LightningTestModel(hparams)
-
-    save_dir = init_save_dir()
+    model = CurrentTestModel(hparams)
 
     # exp file to get meta
     trainer_options = dict(
         max_nb_epochs=1,
         val_percent_check=0.1,
-        train_percent_check=0.1,
+        train_percent_check=1.0,
     )
 
     # fit model
@@ -945,6 +1266,40 @@ def test_multiple_val_dataloader():
     [run_prediction(dataloader, trainer.model) for dataloader in trainer.val_dataloader]
 
 
+def test_multiple_test_dataloader():
+    """
+    Verify multiple test_dataloader
+    :return:
+    """
+    class CurrentTestModel(
+        LightningTestMultipleDataloadersMixin,
+        LightningTestModelBase
+    ):
+        pass
+    hparams = get_hparams()
+    model = CurrentTestModel(hparams)
+
+    # exp file to get meta
+    trainer_options = dict(
+        max_nb_epochs=1,
+        val_percent_check=0.1,
+        train_percent_check=0.1,
+    )
+
+    # fit model
+    trainer = Trainer(**trainer_options)
+    result = trainer.fit(model)
+
+    # verify there are 2 val loaders
+    assert len(trainer.test_dataloader) == 2, 'Multiple test_dataloaders not initiated properly'
+
+    # make sure predictions are good for each test set
+    [run_prediction(dataloader, trainer.model) for dataloader in trainer.test_dataloader]
+
+    # run the test method
+    trainer.test()
+
+
 # ------------------------------------------------------------------------
 # UTILS
 # ------------------------------------------------------------------------
@@ -973,7 +1328,7 @@ def run_gpu_model_test(trainer_options, model, hparams, on_gpu=True):
     # test model loading
     pretrained_model = load_model(exp, save_dir, on_gpu)
 
-    # test model preds
+    # test new model accuracy
     run_prediction(model.test_dataloader, pretrained_model)
 
     if trainer.use_ddp:
@@ -1024,7 +1379,8 @@ def get_model(use_test_model=False):
 def get_exp(debug=True, version=None):
     # set up exp object without actually saving logs
     root_dir = os.path.dirname(os.path.realpath(__file__))
-    exp = Experiment(debug=debug, save_dir=root_dir, name='tests_tt_dir', version=version)
+    save_dir = os.path.join(root_dir, 'save_dir')
+    exp = Experiment(debug=debug, save_dir=save_dir, name='tests_tt_dir', version=version)
     return exp
 
 
@@ -1033,7 +1389,8 @@ def init_save_dir():
     save_dir = os.path.join(root_dir, 'save_dir')
 
     if os.path.exists(save_dir):
-        shutil.rmtree(save_dir)
+        n = np.random.randint(0, 10000000, 1)[0]
+        shutil.move(save_dir, save_dir + f'_{n}')
 
     os.makedirs(save_dir, exist_ok=True)
 
@@ -1044,10 +1401,11 @@ def clear_save_dir():
     root_dir = os.path.dirname(os.path.realpath(__file__))
     save_dir = os.path.join(root_dir, 'save_dir')
     if os.path.exists(save_dir):
-        shutil.rmtree(save_dir)
+        n = np.random.randint(0, 10000000, 1)[0]
+        shutil.move(save_dir, save_dir + f'_{n}')
 
 
-def load_model(exp, save_dir, on_gpu, map_location=None):
+def load_model(exp, save_dir, on_gpu, module_class=LightningTemplateModel):
 
     # load trained model
     tags_path = exp.get_data_path(exp.name, exp.version)
@@ -1056,17 +1414,16 @@ def load_model(exp, save_dir, on_gpu, map_location=None):
     checkpoints = [x for x in os.listdir(save_dir) if '.ckpt' in x]
     weights_dir = os.path.join(save_dir, checkpoints[0])
 
-    trained_model = LightningTemplateModel.load_from_metrics(weights_path=weights_dir,
-                                                             tags_csv=tags_path,
-                                                             on_gpu=on_gpu,
-                                                             map_location=map_location)
+    trained_model = module_class.load_from_metrics(weights_path=weights_dir,
+                                                   tags_csv=tags_path,
+                                                   on_gpu=on_gpu)
 
     assert trained_model is not None, 'loading model failed'
 
     return trained_model
 
 
-def run_prediction(dataloader, trained_model):
+def run_prediction(dataloader, trained_model, dp=False):
     # run prediction on 1 batch
     for batch in dataloader:
         break
@@ -1074,23 +1431,45 @@ def run_prediction(dataloader, trained_model):
     x, y = batch
     x = x.view(x.size(0), -1)
 
-    y_hat = trained_model(x)
+    if dp:
+        output = trained_model(batch, 0)
+        acc = output['val_acc']
+        acc = torch.mean(acc).item()
 
-    # acc
-    labels_hat = torch.argmax(y_hat, dim=1)
-    val_acc = torch.sum(y == labels_hat).item() / (len(y) * 1.0)
-    val_acc = torch.tensor(val_acc)
-    val_acc = val_acc.item()
+    else:
+        y_hat = trained_model(x)
 
-    print(val_acc)
+        # acc
+        labels_hat = torch.argmax(y_hat, dim=1)
+        acc = torch.sum(y == labels_hat).item() / (len(y) * 1.0)
+        acc = torch.tensor(acc)
+        acc = acc.item()
 
-    assert val_acc > 0.50, 'this model is expected to get > 0.50 in test set (it got %f)' % val_acc
+    assert acc > 0.50, f'this model is expected to get > 0.50 in test set (it got {acc})'
 
 
-def assert_ok_acc(trainer):
+def assert_ok_val_acc(trainer):
     # this model should get 0.80+ acc
     acc = trainer.tng_tqdm_dic['val_acc']
-    assert acc > 0.50, 'model failed to get expected 0.50 validation accuracy. Got: %f' % acc
+    assert acc > 0.50, f'model failed to get expected 0.50 validation accuracy. Got: {acc}'
+
+
+def assert_ok_test_acc(trainer):
+    # this model should get 0.80+ acc
+    acc = trainer.tng_tqdm_dic['test_acc']
+    assert acc > 0.50, f'model failed to get expected 0.50 validation accuracy. Got: {acc}'
+
+
+def can_run_gpu_test():
+    if not torch.cuda.is_available():
+        warnings.warn('test_multi_gpu_model_ddp cannot run.'
+                      ' Rerun on a GPU node to run this test')
+        return False
+    if not torch.cuda.device_count() > 1:
+        warnings.warn('test_multi_gpu_model_ddp cannot run.'
+                      ' Rerun on a node with 2+ GPUs to run this test')
+        return False
+    return True
 
 
 if __name__ == '__main__':

tox.ini

@@ -34,11 +34,11 @@ deps =
 commands =
     check-manifest --ignore tox.ini
     python setup.py check -m -s
-    coverage run --source pytorch_lightning -m py.test pytorch_lightning tests examples -v --doctest-modules
     flake8 .
+    coverage run --source pytorch_lightning -m py.test pytorch_lightning tests examples -v --doctest-modules
 
 [flake8]
-exclude = .tox,*.egg,build,temp,examples/templates
+exclude = .tox,*.egg,build,temp,examples/*
 select = E,W,F
 doctests = True
 verbose = 2