release v0.5.0

Rename variables (#124 )
- data_batch → batch - batch_i → batch_idx - dataloader_i → dataloader_idx - tng → training - training_dataloader → train_dataloader - add_log_row_interval → row_log_interval - gradient_clip → gradient_clip_val - prog → progress - tqdm_dic → tqdm_dict
2026-09-11 12:31:23 +08:00 · 2019-09-26 10:42:24 -04:00 · 2019-09-25 19:05:06 -04:00 · 2019-09-25 14:52:40 -04:00 · 2019-09-21 09:45:12 -04:00 · 2019-09-21 09:44:08 -04:00
54 changed files with 2814 additions and 1371 deletions
@@ -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 🙃
@@ -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
@@ -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
@@ -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.    

@@ -108,7 +110,7 @@ class CoolSystem(pl.LightningModule):
        return torch.optim.Adam(self.parameters(), lr=0.02)

    @pl.data_loader
-    def tng_dataloader(self):
+    def train_dataloader(self):
        # REQUIRED
        return DataLoader(MNIST(os.getcwd(), train=True, download=True, transform=transforms.ToTensor()), batch_size=32)

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

@@ -145,11 +146,14 @@ exp = Experiment(save_dir=os.getcwd())
 # pass in experiment for automatic tensorboard logging.    
 trainer = Trainer(experiment=exp, max_nb_epochs=1, train_percent_check=0.1)

-# train on 4 gpus
-# trainer = Trainer(experiment=exp, max_nb_epochs=1, gpus=[0, 1, 2, 3])
+# train on 4 gpus (lightning chooses GPUs for you)
+# trainer = Trainer(experiment=exp, max_nb_epochs=1, gpus=4)  
+
+# train on 4 gpus (you choose GPUs)
+# trainer = Trainer(experiment=exp, max_nb_epochs=1, gpus=[0, 1, 3, 7])   

 # train on 32 gpus across 4 nodes (make sure to submit appropriate SLURM job)
-# trainer = Trainer(experiment=exp, max_nb_epochs=1, gpus=[0, 1, 2, 3, 4, 5, 6, 7], nb_gpu_nodes=4)
+# trainer = Trainer(experiment=exp, max_nb_epochs=1, gpus=8, nb_gpu_nodes=4)

 # train (1 epoch only here for demo)
 trainer.fit(model)
@@ -159,6 +163,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!    
@@ -168,16 +177,16 @@ You define the blue parts using the LightningModule interface:

 ```python
 # what to do in the training loop
-def training_step(self, data_batch, batch_nb):
+def training_step(self, batch, batch_nb):

 # what to do in the validation loop
-def validation_step(self, data_batch, batch_nb):
+def validation_step(self, batch, batch_nb):

 # how to aggregate validation_step outputs
 def validation_end(self, outputs):

 # and your dataloaders
-def tng_dataloader():
+def train_dataloader():
 def val_dataloader():
 def test_dataloader():
 ```
@@ -186,8 +195,8 @@ def test_dataloader():

 ```python
 # define what happens for training here
-def training_step(self, data_batch, batch_nb):
-    x, y = data_batch
+def training_step(self, batch, batch_nb):
+    x, y = batch
    
    # define your own forward and loss calculation
    hidden_states = self.encoder(x)
@@ -213,8 +222,8 @@ def training_step(self, data_batch, batch_nb):

 ```python
 # define what happens for validation here
-def validation_step(self, data_batch, batch_nb):    
-    x, y = data_batch
+def validation_step(self, batch, batch_nb):    
+    x, y = batch
    
    # or as basic as a CNN classification
    out = self.forward(x)
@@ -239,8 +248,8 @@ def validation_end(self, outputs):

    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
+    tqdm_dict = {'val_loss': val_loss_mean.item(), 'val_acc': val_acc_mean.item()}
+    return tqdm_dict
 ```
   
 ## Tensorboard    
@@ -252,10 +261,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 +279,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 +301,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 +310,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,10 +320,11 @@ 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)
+- [Early stopping](https://williamfalcon.github.io/pytorch-lightning/Trainer/Training%20Loop/#early-stopping)
 - [Force disable early stop](https://williamfalcon.github.io/pytorch-lightning/Trainer/Training%20Loop/#force-disable-early-stop)
 - [Gradient Clipping](https://williamfalcon.github.io/pytorch-lightning/Trainer/Training%20Loop/#gradient-clipping)
 - [Hooks](https://williamfalcon.github.io/pytorch-lightning/Trainer/hooks/)
@@ -323,7 +333,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 +342,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 +383,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 +409,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
 ```
@@ -10,18 +10,20 @@ Otherwise, to Define a Lightning Module, implement the following methods:
 **Required**:  

 - [training_step](RequiredTrainerInterface.md#training_step)      
- [tng_dataloader](RequiredTrainerInterface.md#tng_dataloader)    
+- [train_dataloader](RequiredTrainerInterface.md#train_dataloader)    
 - [configure_optimizers](RequiredTrainerInterface.md#configure_optimizers)    

 **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)    
 - [on_load_checkpoint](RequiredTrainerInterface.md#on_load_checkpoint)    
- [update_tng_log_metrics](RequiredTrainerInterface.md#update_tng_log_metrics)    
+- [update_training_log_metrics](RequiredTrainerInterface.md#update_training_log_metrics)    
 - [add_model_specific_args](RequiredTrainerInterface.md#add_model_specific_args)    

 ---
@@ -63,12 +65,23 @@ 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)]

    @pl.data_loader
-    def tng_dataloader(self):
+    def train_dataloader(self):
        return DataLoader(MNIST(os.getcwd(), train=True, download=True, transform=transforms.ToTensor()), batch_size=32)

    @pl.data_loader
@@ -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)
 ```
 ---   
@@ -97,7 +111,7 @@ The LightningModule interface is on the right. Each method corresponds to a part
 ### training_step

 ``` {.python}
-def training_step(self, data_batch, batch_nb)
+def training_step(self, batch, batch_nb)
 ```

 In this step you'd normally do the forward pass and calculate the loss for a batch. You can also do fancier things like multiple forward passes or something specific to your model.
@@ -106,7 +120,7 @@ In this step you'd normally do the forward pass and calculate the loss for a bat

 | Param  | description  |
 |---|---|
-|  data_batch | The output of your dataloader. A tensor, tuple or list  |
+|  batch | The output of your dataloader. A tensor, tuple or list  |
 |  batch_nb | Integer displaying which batch this is  |

 **Return**   
@@ -116,14 +130,14 @@ Dictionary or OrderedDict
 | key  | value  | is required |
 |---|---|---|
 |  loss | tensor scalar  | Y |
-|  prog | Dict for progress bar display. Must have only tensors  | N |
+|  progress | Dict for progress bar display. Must have only tensors  | N |


 **Example**

 ``` {.python}
-def training_step(self, data_batch, batch_nb):
-    x, y, z = data_batch
+def training_step(self, batch, batch_nb):
+    x, y, z = batch
    
    # implement your own
    out = self.forward(x)
@@ -131,7 +145,7 @@ def training_step(self, data_batch, batch_nb):
    
    output = {
        'loss': loss, # required
-        'prog': {'tng_loss': loss, 'batch_nb': batch_nb} # optional
+        'progress': {'training_loss': loss, 'batch_nb': batch_nb} # optional
    }
    
    # return a dict
@@ -141,7 +155,7 @@ def training_step(self, data_batch, batch_nb):
 If you define multiple optimizers, this step will also be called with an additional ```optimizer_idx``` param.    
 ``` {.python}
 # Multiple optimizers (ie: GANs)     
-def training_step(self, data_batch, batch_nb, optimizer_idx):
+def training_step(self, batch, batch_nb, optimizer_idx):
    if optimizer_idx == 0:
        # do training_step with encoder
    if optimizer_idx == 1:
@@ -149,11 +163,11 @@ def training_step(self, data_batch, batch_nb, optimizer_idx):
 ```    

 --- 
-### tng_dataloader 
+### train_dataloader 

 ``` {.python}
@pl.data_loader
-def tng_dataloader(self)
+def train_dataloader(self)
 ```
 Called by lightning during training loop. Make sure to use the @pl.data_loader decorator, this ensures not calling this function until the data are needed.

@@ -164,7 +178,7 @@ PyTorch DataLoader

 ``` {.python}
@pl.data_loader
-def tng_dataloader(self):
+def train_dataloader(self):
    transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,), (1.0,))])
    dataset = MNIST(root='/path/to/mnist/', train=True, transform=transform, download=True)
    loader = torch.utils.data.DataLoader(
@@ -225,25 +239,26 @@ the [optimizer_step](https://williamfalcon.github.io/pytorch-lightning/Trainer/h
 ### validation_step

 ``` {.python}
-def validation_step(self, data_batch, batch_nb)   
+# if you have one val dataloader:
+def validation_step(self, batch, batch_nb)   

-# if have multiple val dataloaders:  
-def validation_step(self, data_batch, batch_nb, dataloader_idx)
+# if you have multiple val dataloaders:  
+def validation_step(self, batch, batch_nb, dataloader_idxdx)
 ```
 **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**   

 | Param  | description  |
 |---|---|
-|  data_batch | The output of your dataloader. A tensor, tuple or list  |
+|  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 val datasets used)  |
+|  dataloader_idx | Integer displaying which dataloader this is (only if multiple val datasets used)  |

 **Return**   

@@ -255,12 +270,12 @@ 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
+def validation_step(self, batch, batch_nb):
+    x, y = 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
@@ -287,7 +302,7 @@ If you pass in multiple validation datasets, validation_step will have an additi

 ```python
 # CASE 2: multiple validation datasets
-def validation_step(self, data_batch, batch_nb, dataset_idx):
+def validation_step(self, batch, batch_nb, dataset_idx):
    # dataset_idx tells you which dataset this is.   
 ```   

@@ -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):
    """
@@ -334,8 +351,173 @@ def validation_end(self, outputs):

    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
+    tqdm_dict = {'val_loss': val_loss_mean.item(), 'val_acc': val_acc_mean.item()}
+    return tqdm_dict
+```
+
+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_dict = {'val_loss': val_loss_mean.item(), 'val_acc': val_acc_mean.item()}
+    return tqdm_dict
+```
+
+### test_step
+
+``` {.python}
+# if you have one test dataloader:
+def test_step(self, batch, batch_nb)   
+
+# if you have multiple test dataloaders:  
+def test_step(self, batch, batch_nb, dataloader_idxdx)
+```
+**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  |
+|---|---|
+|  batch | The output of your dataloader. A tensor, tuple or list  |
+|  batch_nb | Integer displaying which batch this is  |
+|  dataloader_idx | 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, batch, batch_nb):
+    x, y = 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, 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_dict = {'test_loss': test_loss_mean.item(), 'test_acc': test_acc_mean.item()}
+    return tqdm_dict
+```
+
+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_dict = {'test_loss': test_loss_mean.item(), 'test_acc': test_acc_mean.item()}
+    return tqdm_dict
 ```

 --- 
@@ -451,10 +633,10 @@ def test_dataloader(self):
 ```

 --- 
-### update_tng_log_metrics 
+### update_training_log_metrics 

 ``` {.python}
-def update_tng_log_metrics(self, logs)
+def update_training_log_metrics(self, logs)
 ```
 Called by lightning right before it logs metrics for this batch.
 This is a chance to amend or add to the metrics about to be logged.
@@ -465,7 +647,7 @@ Dict
 **Example**

 ``` {.python}
-def update_tng_log_metrics(self, logs):
+def update_training_log_metrics(self, logs):
    # modify or add to logs
    return logs
 ```
@@ -492,7 +674,7 @@ def add_model_specific_args(parent_parser, root_dir):
    parser = HyperOptArgumentParser(strategy=parent_parser.strategy, parents=[parent_parser])

    # param overwrites
-    # parser.set_defaults(gradient_clip=5.0)
+    # parser.set_defaults(gradient_clip_val=5.0)

    # network params
    parser.opt_list('--drop_prob', default=0.2, options=[0.2, 0.5], type=float, tunable=False)
@@ -22,7 +22,7 @@ self.experiment.add_scalars(...)
 Total training batches seen across all epochs   

 --- 
-#### gradient_clip   
+#### gradient_clip_val   
 The current gradient clip value    

 ---
@@ -8,11 +8,12 @@ To enable checkpointing, define the checkpoint callback and give it to the train
 from pytorch_lightning.callbacks import ModelCheckpoint

 checkpoint_callback = ModelCheckpoint(
-    filepath='/path/to/store/weights.ckpt',
+    filepath='/path/to/store/weights/',
    save_best_only=True,
    verbose=True,
    monitor='val_loss',
-    mode='min'
+    mode='min',
+    prefix=''
 )

 trainer = Trainer(checkpoint_callback=checkpoint_callback)
@@ -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, distributed_backend='dp')``` |
+|   | 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, distributed_backend='ddp')
 ```

 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.
@@ -122,6 +140,9 @@ cluster.per_experiment_nb_gpus = 8
 cluster.add_slurm_cmd(cmd='ntasks-per-node', value=8, comment='1 task per gpu')
 ```

+**NOTE:** When running in DDP mode, any errors in your code will show up as an NCCL issue.
+Set the ```NCCL_DEBUG=INFO``` flag to see the ACTUAL error.
+
 Finally, make sure to add a distributed sampler to your dataset. The distributed sampler copies a 
 portion of your dataset onto each GPU. (World_size = gpus_per_node * nb_nodes).   

@@ -13,7 +13,15 @@ trainer = Trainer(show_progress_bar=True)
 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)
+trainer = Trainer(row_log_interval=10)
+```   
+
+---
+#### Log metric row every k batches 
+Logs GPU memory when metrics are logged.   
+``` {.python}
+# DEFAULT
+trainer = Trainer(log_gpu_memory=False)
 ```

 ---
@@ -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
@@ -23,6 +25,9 @@ parser.opt_list('--nb_layers', default=2, type=int, tunable=True, options=[2, 4,
 hparams = parser.parse_args()    
 ```    
    
+**NOTE** You must set ```Tunable=True``` for that argument to be considered in the permutation set. Otherwise
+test-tube will use the default value. This flag is useful when you don't want to search over an argument and
+want to use the default instead.   
     
 (2). Define the cluster options in the [SlurmCluster object](https://williamfalcon.github.io/test-tube/hpc/SlurmCluster/) (over 5 nodes and 8 gpus)    

@@ -55,8 +60,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 +71,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 +86,27 @@ 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!
-
+**NOTE** nb_trials specifies how many of the possible permutations to use. If using ```grid_search``` it will use
+the depth first ordering. If using ```random_search``` it will use the first k shuffled options. FYI, random search
+has been shown to be just as good as any Bayesian optimization method when using a reasonable number of samples (60),
+[see this paper for more information](http://www.jmlr.org/papers/volume13/bergstra12a/bergstra12a.pdf).

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

@@ -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...)  
+ 
@@ -19,6 +19,24 @@ It can be useful to force training for a minimum number of epochs or limit to a
 trainer = Trainer(min_nb_epochs=1, max_nb_epochs=1000)
 ```

+---
+#### Early stopping
+To enable ealry-stopping, define the callback and give it to the trainer.
+``` {.python}
+from pytorch_lightning.callbacks import EarlyStopping
+
+# DEFAULTS
+early_stop_callback = EarlyStopping(
+    monitor='val_loss',
+    min_delta=0.00,
+    patience=0,
+    verbose=False,
+    mode='auto'
+)
+
+trainer = Trainer(early_stop_callback=early_stop_callback)
+```
+
 ---
 #### Force disable early stop 
 Use this to turn off early stopping and run training to the [max_epoch](#force-training-for-min-or-max-epochs)
@@ -34,10 +52,10 @@ Specifically, this will [clip the gradient norm computed over all model paramete

 ``` {.python}
 # DEFAULT (ie: don't clip)
-trainer = Trainer(gradient_clip=0)
+trainer = Trainer(gradient_clip_val=0)

 # clip gradients with norm above 0.5
-trainer = Trainer(gradient_clip=0.5)
+trainer = Trainer(gradient_clip_val=0.5)
 ```

 ---
@@ -5,8 +5,6 @@ Below are all the things lightning automates for you in the validation loop.
 Lightning will run 5 steps of validation in the beginning of training as a sanity check so you don't have to wait until a full epoch to catch possible validation issues.


-
-
 ---
 #### Check validation every n epochs
 If you have a small dataset you might want to check validation every n epochs
@@ -60,4 +58,6 @@ Lightning runs a few steps of validation in the beginning of training. This avoi
 ``` {.python}
 # DEFAULT
 trainer = Trainer(nb_sanity_val_steps=5)
-```
+```
+
+You can use `Trainer(nb_sanity_val_steps=0)` to skip the sanity check.
@@ -58,12 +58,12 @@ def on_post_performance_check(self):
 ```

 ---
-#### on_tng_metrics
+#### on_training_metrics
 Called in the training loop, right before metrics are logged.
 Although you can log at any time by using self.experiment, you can use
 this callback to modify what will be logged.
 ```python
-def on_tng_metrics(self, metrics):
+def on_training_metrics(self, metrics):
    # do something before validation end
 ```

@@ -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/)  
@@ -40,7 +40,7 @@ The main function should have 3 arguments:
 - slurm_manager: Slurm cluster manager object (can be None)
 - dict: for you to return any values you want (useful in meta-learning, otherwise set to _)    

-```{}
+```python
 def main(hparams, cluster, results_dict):
    """
    Main training routine specific for this project
@@ -120,6 +120,7 @@ Notice a few things about this flow:

 - [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)
+- [Early stopping](https://williamfalcon.github.io/pytorch-lightning/Trainer/Training%20Loop/#early-stopping)
 - [Force disable early stop](https://williamfalcon.github.io/pytorch-lightning/Trainer/Training%20Loop/#force-disable-early-stop)
 - [Gradient Clipping](https://williamfalcon.github.io/pytorch-lightning/Trainer/Training%20Loop/#gradient-clipping)
 - [Hooks](https://williamfalcon.github.io/pytorch-lightning/Trainer/hooks/)
@@ -137,3 +138,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/)  
@@ -79,14 +79,14 @@ class LightningTemplateModel(LightningModule):
        nll = F.nll_loss(logits, labels)
        return nll

-    def training_step(self, data_batch, batch_i):
+    def training_step(self, batch, batch_idx):
        """
        Lightning calls this inside the training loop
-        :param data_batch:
+        :param batch:
        :return:
        """
        # forward pass
-        x, y = data_batch
+        x, y = batch
        x = x.view(x.size(0), -1)

        y_hat = self.forward(x)
@@ -105,13 +105,13 @@ class LightningTemplateModel(LightningModule):
        # can also return just a scalar instead of a dict (return loss_val)
        return output

-    def validation_step(self, data_batch, batch_i):
+    def validation_step(self, batch, batch_idx):
        """
        Lightning calls this inside the validation loop
-        :param data_batch:
+        :param batch:
        :return:
        """
-        x, y = data_batch
+        x, y = batch
        x = x.view(x.size(0), -1)
        y_hat = self.forward(x)

@@ -167,8 +167,8 @@ class LightningTemplateModel(LightningModule):

        val_loss_mean /= len(outputs)
        val_acc_mean /= len(outputs)
-        tqdm_dic = {'val_loss': val_loss_mean, 'val_acc': val_acc_mean}
-        return tqdm_dic
+        tqdm_dict = {'val_loss': val_loss_mean, 'val_acc': val_acc_mean}
+        return tqdm_dict

    # ---------------------
    # TRAINING SETUP
@@ -208,8 +208,8 @@ class LightningTemplateModel(LightningModule):
        return loader

    @pl.data_loader
-    def tng_dataloader(self):
-        print('tng data loader called')
+    def train_dataloader(self):
+        print('training data loader called')
        return self.__dataloader(train=True)

    @pl.data_loader
@@ -233,22 +233,22 @@ class LightningTemplateModel(LightningModule):
        parser = HyperOptArgumentParser(strategy=parent_parser.strategy, parents=[parent_parser])

        # param overwrites
-        # parser.set_defaults(gradient_clip=5.0)
+        # parser.set_defaults(gradient_clip_val=5.0)

        # network params
        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)
-        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)

@@ -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.   
@@ -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
@@ -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()
@@ -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
@@ -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()
@@ -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)
@@ -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,
@@ -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
@@ -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
@@ -71,7 +71,7 @@ def main(hparams):
        checkpoint_callback=checkpoint,
        early_stop_callback=early_stop,
        gpus=hparams.gpus,
-        distributed_backend='ddp'
+        distributed_backend=hparams.dist_backend
    )

    # ------------------------
@@ -95,6 +95,8 @@ if __name__ == '__main__':
    parent_parser.add_argument('--gpus', type=str, default='-1',
                               help='how many gpus to use in the node.'
                                    ' value -1 uses all the gpus on the node')
+    parent_parser.add_argument('--dist_backend', type=str, default='ddp',
+                                help='When using multiple GPUs set Trainer(distributed_backend=dp) (or ddp)')     
    parent_parser.add_argument('--test_tube_save_path', type=str, default=test_tube_dir,
                               help='where to save logs')
    parent_parser.add_argument('--model_save_path', type=str, default=checkpoint_dir,
@@ -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
@@ -71,6 +71,7 @@ def main(hparams):
        checkpoint_callback=checkpoint,
        early_stop_callback=early_stop,
        gpus=hparams.gpus,
+        distributed_backend=hparams.dist_backend,
    )

    # ------------------------
@@ -94,6 +95,8 @@ if __name__ == '__main__':
    parent_parser.add_argument('--gpus', type=str, default='-1',
                               help='how many gpus to use in the node.'
                                    ' value -1 uses all the gpus on the node')
+    parent_parser.add_argument('--dist_backend', type=str, default='dp',
+                               help='When using multiple GPUs set Trainer(distributed_backend=dp) (or ddp)')
    parent_parser.add_argument('--test_tube_save_path', type=str, default=test_tube_dir,
                               help='where to save logs')
    parent_parser.add_argument('--model_save_path', type=str, default=checkpoint_dir,
@@ -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

@@ -146,7 +146,7 @@ class GAN(pl.LightningModule):
        return [opt_g, opt_d], []

    @pl.data_loader
-    def tng_dataloader(self):
+    def train_dataloader(self):
        transform = transforms.Compose([transforms.ToTensor(),
                                        transforms.Normalize([0.5], [0.5])])
        dataset = MNIST(os.getcwd(), train=True, download=True, transform=transform)
@@ -2,6 +2,7 @@ site_name: PyTorch lightning Documentation
 theme:
  name: 'material'
 docs_dir: docs
+repo_name: 'williamFalcon/pytorch-lightning'
 repo_url: https://github.com/williamFalcon/pytorch-lightning
 site_dir: 'site'
 site_description: 'Documentation for PyTorch LightningModule, the researcher version of keras.'
@@ -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

@@ -1,6 +1,7 @@
-from .pt_callbacks import EarlyStopping, ModelCheckpoint
+from .pt_callbacks import EarlyStopping, ModelCheckpoint, GradientAccumulationScheduler

 __all__ = [
    'EarlyStopping',
    'ModelCheckpoint',
+    'GradientAccumulationScheduler',
 ]
@@ -1,5 +1,6 @@
 import os
 import shutil
+import warnings

 import numpy as np

@@ -11,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
@@ -254,6 +254,37 @@ class ModelCheckpoint(Callback):
                self.save_model(filepath, overwrite=False)


+class GradientAccumulationScheduler(Callback):
+    """Change gradient accumulation factor according to scheduling.
+    # Arguments
+        scheduling: dict, scheduling in format {epoch: accumulation_factor}
+    """
+    def __init__(self, scheduling: dict):
+        if scheduling == {}:  # empty dict error
+            raise TypeError("Empty dict cannot be interpreted correct")
+
+        for key in scheduling.keys():
+            if not isinstance(key, int) or not isinstance(scheduling[key], int):
+                raise TypeError("All epoches and accumulation factor must be integers")
+
+        minimal_epoch = min(scheduling.keys())
+        if minimal_epoch < 1:
+            msg = f"Epochs indexing from 1, epoch {minimal_epoch} cannot be interpreted correct"
+            raise IndexError(msg)
+        elif minimal_epoch != 1:  # if user didnt define first epoch accumulation factor
+            scheduling.update({1: 1})
+
+        self.scheduling = scheduling
+        self.epochs = sorted(scheduling.keys())
+
+    def on_epoch_begin(self, epoch, trainer):
+        epoch += 1  # indexing epochs from 1
+        for i in reversed(range(len(self.epochs))):
+            if epoch >= self.epochs[i]:
+                trainer.accumulate_grad_batches = self.scheduling.get(self.epochs[i])
+                break
+
+
 if __name__ == '__main__':
    c = EarlyStopping(min_delta=0.9, patience=2, verbose=True)
    losses = [10, 9, 8, 8, 6, 4.3, 5, 4.4, 2.8, 2.5]
@@ -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)
                # ---------------
@@ -5,12 +5,12 @@ class ModelHooks(torch.nn.Module):

    def on_sanity_check_start(self):
        """
-        Called before starting validate
+        Called before starting evaluate
        :return:
        """
        pass

-    def on_batch_start(self, data_batch):
+    def on_batch_start(self, batch):
        pass

    def on_batch_end(self):
@@ -28,7 +28,7 @@ class ModelHooks(torch.nn.Module):
    def on_post_performance_check(self):
        pass

-    def on_tng_metrics(self, metrics):
+    def on_training_metrics(self, metrics):
        pass

    def on_before_zero_grad(self, optimizer):
@@ -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
@@ -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)
@@ -86,7 +106,7 @@ class LightningModule(GradInformation, ModelIO, ModelHooks):
        optimizer.zero_grad()

    @data_loader
-    def tng_dataloader(self):
+    def train_dataloader(self):
        """
        Implement a PyTorch DataLoader
        :return:
@@ -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)
@@ -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,
+)
@@ -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):
+    def on_training_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
@@ -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
@@ -80,14 +81,14 @@ class NoValModel(LightningModule):
        nll = F.nll_loss(logits, labels)
        return nll

-    def training_step(self, data_batch, batch_i):
+    def training_step(self, batch, batch_idx):
        """
        Lightning calls this inside the training loop
-        :param data_batch:
+        :param batch:
        :return:
        """
        # forward pass
-        x, y = data_batch
+        x, y = batch
        x = x.view(x.size(0), -1)

        y_hat = self.forward(x)
@@ -103,15 +104,12 @@ class NoValModel(LightningModule):
        if self.trainer.batch_nb % 1 == 0:
            output = OrderedDict({
                'loss': loss_val,
-                'prog': {'some_val': loss_val * loss_val}
+                'progress': {'some_val': loss_val * loss_val}
            })
            return output
        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,))])
@@ -155,8 +153,8 @@ class NoValModel(LightningModule):
        return loader

    @data_loader
-    def tng_dataloader(self):
-        return self.__dataloader(train=True)
+    def train_dataloader(self):
+        return self._dataloader(train=True)

    @staticmethod
    def add_model_specific_args(parent_parser, root_dir):  # pragma: no cover
@@ -169,7 +167,7 @@ class NoValModel(LightningModule):
        parser = HyperOptArgumentParser(strategy=parent_parser.strategy, parents=[parent_parser])

        # param overwrites
-        # parser.set_defaults(gradient_clip=5.0)
+        # parser.set_defaults(gradient_clip_val=5.0)

        # network params
        parser.opt_list('--drop_prob', default=0.2, options=[0.2, 0.5], type=float, tunable=False)
@@ -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, batch, batch_idx):
+        """
+        Lightning calls this inside the validation loop
+        :param batch:
+        :return:
+        """
+        x, y = 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_idx % 1 == 0:
+            output = OrderedDict({
+                'val_loss': loss_val,
+                'val_acc': val_acc,
+            })
+            return output
+        if batch_idx % 2 == 0:
+            return val_acc
+
+        if batch_idx % 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_dict = {'val_loss': val_loss_mean.item(), 'val_acc': val_acc_mean.item()}
+        return tqdm_dict
+
+
+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, batch, batch_idx, dataloader_idx):
+        """
+        Lightning calls this inside the validation loop
+        :param batch:
+        :return:
+        """
+        x, y = 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_idx % 1 == 0:
+            output = OrderedDict({
+                'val_loss': loss_val,
+                'val_acc': val_acc,
+            })
+            return output
+        if batch_idx % 2 == 0:
+            return val_acc
+
+        if batch_idx % 3 == 0:
+            output = OrderedDict({
+                'val_loss': loss_val,
+                'val_acc': val_acc,
+                'test_dic': {'val_loss_a': loss_val}
+            })
+            return output
+        if batch_idx % 5 == 0:
+            output = OrderedDict({
+                f'val_loss_{dataloader_idx}': loss_val,
+                f'val_acc_{dataloader_idx}': 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_dict = {'val_loss': val_loss_mean.item(), 'val_acc': val_acc_mean.item()}
+        return tqdm_dict
+
+
+class LightningTestStepMixin:
+
+    @data_loader
+    def test_dataloader(self):
+        return self._dataloader(train=False)
+
+    def test_step(self, batch, batch_idx):
+        """
+        Lightning calls this inside the validation loop
+        :param batch:
+        :return:
+        """
+        x, y = 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_idx % 1 == 0:
+            output = OrderedDict({
+                'test_loss': loss_test,
+                'test_acc': test_acc,
+            })
+            return output
+        if batch_idx % 2 == 0:
+            return test_acc
+
+        if batch_idx % 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_dict = {'test_loss': test_loss_mean.item(), 'test_acc': test_acc_mean.item()}
+        return tqdm_dict
+
+
+class LightningTestStepMultipleDataloadersMixin:
+
+    @data_loader
+    def test_dataloader(self):
+        return [self._dataloader(train=False), self._dataloader(train=False)]
+
+    def test_step(self, batch, batch_idx, dataloader_idx):
+        """
+        Lightning calls this inside the validation loop
+        :param batch:
+        :return:
+        """
+        x, y = 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_idx % 1 == 0:
+            output = OrderedDict({
+                'test_loss': loss_test,
+                'test_acc': test_acc,
+            })
+            return output
+        if batch_idx % 2 == 0:
+            return test_acc
+
+        if batch_idx % 3 == 0:
+            output = OrderedDict({
+                'test_loss': loss_test,
+                'test_acc': test_acc,
+                'test_dic': {'test_loss_a': loss_test}
+            })
+            return output
+        if batch_idx % 5 == 0:
+            output = OrderedDict({
+                f'test_loss_{dataloader_idx}': loss_test,
+                f'test_acc_{dataloader_idx}': 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_dict = {'test_loss': test_loss_mean.item(), 'test_acc': test_acc_mean.item()}
+        return tqdm_dict
@@ -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
@@ -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
@@ -8,7 +8,7 @@ import os

 def add_default_args(parser, root_dir, rand_seed=None, possible_model_names=None):

-    # tng, test, val check intervals
+    # training, test, val check intervals
    parser.add_argument('--eval_test_set', dest='eval_test_set', action='store_true',
                        help='true = run test set also')
    parser.add_argument('--check_val_every_n_epoch', default=1, type=int,
@@ -19,7 +19,7 @@ def add_default_args(parser, root_dir, rand_seed=None, possible_model_names=None
    parser.add_argument('--max_nb_epochs', default=200, type=int, help='cap epochs')
    parser.add_argument('--min_nb_epochs', default=2, type=int, help='min epochs')
    parser.add_argument('--train_percent_check', default=1.0, type=float,
-                        help='how much of tng set to check')
+                        help='how much of training set to check')
    parser.add_argument('--val_percent_check', default=1.0, type=float,
                        help='how much of val set to check')
    parser.add_argument('--test_percent_check', default=1.0, type=float,
@@ -29,7 +29,7 @@ def add_default_args(parser, root_dir, rand_seed=None, possible_model_names=None
                        help='how much within 1 epoch to check val')
    parser.add_argument('--log_save_interval', default=100, type=int,
                        help='how many batches between log saves')
-    parser.add_argument('--add_log_row_interval', default=100, type=int,
+    parser.add_argument('--row_log_interval', default=100, type=int,
                        help='add log every k batches')

    # early stopping
@@ -40,7 +40,7 @@ def add_default_args(parser, root_dir, rand_seed=None, possible_model_names=None
                        help='number of epochs until stop')

    # gradient handling
-    parser.add_argument('--gradient_clip', default=-1, type=int)
+    parser.add_argument('--gradient_clip_val', default=-1, type=int)
    parser.add_argument('--track_grad_norm', default=-1, type=int,
                        help='if > 0, will track this grad norm')

@@ -78,9 +78,9 @@ def add_default_args(parser, root_dir, rand_seed=None, possible_model_names=None
    # FAST training
    # use these settings to make sure network has no bugs without running a full dataset
    parser.add_argument('--fast_dev_run', dest='fast_dev_run', default=False, action='store_true',
-                        help='runs validation after 1 tng step')
+                        help='runs validation after 1 training step')
    parser.add_argument('--enable_tqdm', dest='enable_tqdm', default=False, action='store_true',
-                        help='false removes the prog bar')
+                        help='false removes the progress bar')
    parser.add_argument('--overfit', default=-1, type=float,
                        help='% of dataset to use with this option. float, or -1 for none')

@@ -93,7 +93,7 @@ def add_default_args(parser, root_dir, rand_seed=None, possible_model_names=None
    parser.add_argument('--debug', dest='debug', action='store_true',
                        help='enables/disables test tube')
    parser.add_argument('--local', dest='local', action='store_true',
-                        help='enables local tng')
+                        help='enables local training')

    # optimizer
    parser.add_argument('--lr_scheduler_milestones', default=None, type=str)
@@ -14,7 +14,7 @@ from setuptools import setup, find_packages
 # engineer specific practices
 setup(
    name='pytorch-lightning',
-    version='0.4.7',
+    version='0.5.0',
    description='The Keras for ML researchers using PyTorch',
    author='William Falcon',
    author_email='waf2107@columbia.edu',
@@ -30,7 +30,7 @@ setup(
    python_requires='>=3.6',
    install_requires=[
        'torch==1.2.0',
-        'tqdm',
+        'tqdm>=4.35.0',
        'test-tube>=0.6.9',
        'pandas>=0.20.3',
    ],
@@ -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):
@@ -30,7 +32,7 @@ class CoolModel(pl.LightningModule):
    def training_step(self, batch, batch_nb):
        x, y = batch
        y_hat = self.forward(x)
-        return {'tng_loss': self.my_loss(y_hat, y)}
+        return {'training_loss': self.my_loss(y_hat, y)}

    def validation_step(self, batch, batch_nb):
        x, y = batch
@@ -45,7 +47,7 @@ class CoolModel(pl.LightningModule):
        return [torch.optim.Adam(self.parameters(), lr=0.02)]

    @pl.data_loader
-    def tng_dataloader(self):
+    def train_dataloader(self):
        return DataLoader(MNIST('path/to/save', train=True), batch_size=32)

    @pl.data_loader
@@ -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.training_tqdm_dict['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.training_tqdm_dict['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__':
@@ -6,16 +6,29 @@ 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.callbacks import ModelCheckpoint, EarlyStopping
+from pytorch_lightning.testing import (
+    LightningTestModel,
+    LightningTestModelBase,
+    LightningValidationMixin,
+    LightningValidationStepMixin,
+    LightningValidationMultipleDataloadersMixin,
+    LightningTestMixin,
+    LightningTestMultipleDataloadersMixin,
+)
+from pytorch_lightning.callbacks import (
+    ModelCheckpoint,
+    EarlyStopping,
+    GradientAccumulationScheduler,
+)
 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
@@ -26,18 +39,381 @@ 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
+    """
+    # test incorrect configs
+    with pytest.raises(IndexError):
+        assert Trainer(accumulate_grad_batches={0: 3, 1: 4, 4: 6})
+        assert Trainer(accumulate_grad_batches={-2: 3})
+
+    with pytest.raises(TypeError):
+        assert Trainer(accumulate_grad_batches={})
+        assert Trainer(accumulate_grad_batches=[[2, 3], [4, 6]])
+        assert Trainer(accumulate_grad_batches={1: 2, 3.: 4})
+        assert Trainer(accumulate_grad_batches={1: 2.5, 3: 5})
+
+    # test optimizer call freq matches scheduler
+    def optimizer_step(self, epoch_nb, batch_nb, optimizer, optimizer_i):
+        # only test the first 12 batches in epoch
+        if batch_nb < 12:
+            if epoch_nb == 0:
+                # reset counter when starting epoch
+                if batch_nb == 0:
+                    self.prev_called_batch_nb = 0
+
+                    # use this opportunity to test once
+                    assert self.trainer.accumulate_grad_batches == 1
+
+                assert batch_nb == self.prev_called_batch_nb
+                self.prev_called_batch_nb += 1
+
+            elif 1 <= epoch_nb <= 2:
+                # reset counter when starting epoch
+                if batch_nb == 1:
+                    self.prev_called_batch_nb = 1
+
+                    # use this opportunity to test once
+                    assert self.trainer.accumulate_grad_batches == 2
+
+                assert batch_nb == self.prev_called_batch_nb
+                self.prev_called_batch_nb += 2
+
+            else:
+                if batch_nb == 3:
+                    self.prev_called_batch_nb = 3
+
+                    # use this opportunity to test once
+                    assert self.trainer.accumulate_grad_batches == 4
+
+                assert batch_nb == self.prev_called_batch_nb
+                self.prev_called_batch_nb += 3
+
+        optimizer.step()
+
+        # clear gradients
+        optimizer.zero_grad()
+
+    hparams = get_hparams()
+    model = LightningTestModel(hparams)
+    schedule = {1: 2, 3: 4}
+
+    trainer = Trainer(accumulate_grad_batches=schedule,
+                      train_percent_check=0.1,
+                      val_percent_check=0.1,
+                      max_nb_epochs=4)
+
+    # for the test
+    trainer.optimizer_step = optimizer_step
+    model.prev_called_batch_nb = 0
+
+    trainer.fit(model)
+
+
 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])
@@ -85,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()
@@ -141,7 +511,7 @@ def test_early_stopping_cpu_model():
    stopping = EarlyStopping(monitor='val_loss')
    trainer_options = dict(
        early_stop_callback=stopping,
-        gradient_clip=1.0,
+        gradient_clip_val=1.0,
        overfit_pct=0.20,
        track_grad_norm=2,
        print_nan_grads=True,
@@ -165,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()

@@ -176,7 +549,6 @@ def test_no_val_module():

    trainer_options = dict(
        max_nb_epochs=1,
-        cluster=SlurmCluster(),
        experiment=exp,
        checkpoint_callback=ModelCheckpoint(save_dir)
    )
@@ -185,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
@@ -208,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()

@@ -220,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)
    )
@@ -294,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
@@ -373,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])
@@ -390,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
    )
@@ -413,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)
    )
@@ -440,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),
    )
@@ -491,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

@@ -536,7 +929,6 @@ def test_model_saving_loading():

    trainer_options = dict(
        max_nb_epochs=1,
-        cluster=SlurmCluster(),
        experiment=exp,
        checkpoint_callback=ModelCheckpoint(save_dir)
    )
@@ -591,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
@@ -644,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)
@@ -713,7 +1098,7 @@ def test_all_features_cpu_model():
    """

    trainer_options = dict(
-        gradient_clip=1.0,
+        gradient_clip_val=1.0,
        overfit_pct=0.20,
        track_grad_norm=2,
        print_nan_grads=True,
@@ -745,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,
@@ -784,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,
@@ -808,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])
@@ -843,10 +1237,47 @@ def test_multiple_val_dataloader():
    Verify multiple val_dataloader
    :return:
    """
+    class CurrentTestModel(
+        LightningValidationMultipleDataloadersMixin,
+        LightningTestModelBase
+    ):
+        pass
    hparams = get_hparams()
-    model = LightningTestModel(hparams)
+    model = CurrentTestModel(hparams)

-    save_dir = init_save_dir()
+    # exp file to get meta
+    trainer_options = dict(
+        max_nb_epochs=1,
+        val_percent_check=0.1,
+        train_percent_check=1.0,
+    )
+
+    # fit model
+    trainer = Trainer(**trainer_options)
+    result = trainer.fit(model)
+
+    # verify training completed
+    assert result == 1
+
+    # verify there are 2 val loaders
+    assert len(trainer.val_dataloader) == 2, 'Multiple val_dataloaders not initiated properly'
+
+    # make sure predictions are good for each val set
+    [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(
@@ -859,14 +1290,14 @@ def test_multiple_val_dataloader():
    trainer = Trainer(**trainer_options)
    result = trainer.fit(model)

-    # verify tng completed
-    assert result == 1
-
    # verify there are 2 val loaders
-    assert len(trainer.val_dataloader) == 2, 'Multiple val_dataloaders not initiated properly'
+    assert len(trainer.test_dataloader) == 2, 'Multiple test_dataloaders not initiated properly'

-    # make sure predictions are good for each val set
-    [run_prediction(dataloader, trainer.model) for dataloader in trainer.val_dataloader]
+    # 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()


 # ------------------------------------------------------------------------
@@ -897,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:
@@ -948,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


@@ -957,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)

@@ -968,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)
@@ -980,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
@@ -998,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
+    acc = trainer.training_tqdm_dict['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.training_tqdm_dict['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__':
@@ -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
Author	SHA1	Message	Date
William Falcon	c2a0846011	release v0.5.0	2019-09-26 10:42:24 -04:00
Alok Singh	b0a0a47a0b	Rename variables (#124 ) - data_batch → batch - batch_i → batch_idx - dataloader_i → dataloader_idx - tng → training - training_dataloader → train_dataloader - add_log_row_interval → row_log_interval - gradient_clip → gradient_clip_val - prog → progress - tqdm_dic → tqdm_dict	2019-09-25 19:05:06 -04:00
Cola	3d16a686b3	Add EarlyStop documentation (#245 ) * Update Training Loop.md * Update index.md * Update README.md * Update Training Loop.md * Update Training Loop.md	2019-09-25 14:52:40 -04:00
Oscar A. Rangel	eb268c4184	Added missing parameters (#237 ) * Added missing parameters added missing distributed_backend parameter and added the parameter to step 4 Init Trainer. * Update single_gpu_node_dp_template.py	2019-09-21 09:45:12 -04:00
Oscar A. Rangel	6803018a49	changed hard coded paramater, and moved it to parent_parser (#238 ) * changed hard coded paramater, and moved it to parent_parser ```python # ------------------------ # 4 INIT TRAINER # ------------------------ trainer = Trainer( experiment=exp, checkpoint_callback=checkpoint, early_stop_callback=early_stop, gpus=hparams.gpus, distributed_backend=hparams.dist_bak_end ) parent_parser.add_argument('--dist_bak_end', type=str, default='ddp', help='When using multiple GPUs set Trainer(distributed_backend=dp) (or ddp)') ``` * Update single_gpu_node_ddp_template.py	2019-09-21 09:44:08 -04:00
William Falcon	87708157bc	Update trainer.py (#233 )	2019-09-19 08:23:48 -04:00
William Falcon	2a1bc22f42	updated docs	2019-09-17 09:57:16 -04:00
William Falcon	d3afc8acd5	updated docs	2019-09-17 09:53:31 -04:00
William Falcon	4c61d1f30a	updated docs	2019-09-16 11:07:16 -04:00
William Falcon	e1adbe80f9	updated docs	2019-09-16 11:04:40 -04:00
William Falcon	286625a02f	updated docs	2019-09-16 11:02:04 -04:00
William Falcon	b354988255	updated docs	2019-09-16 10:59:28 -04:00
William Falcon	b3c1911813	Update README.md	2019-09-16 10:56:37 -04:00
William Falcon	974afba2be	release v0.4.9	2019-09-16 10:50:59 -04:00
William Falcon	55e7322747	Metrics load (#228 ) * load from metrics defaults to CPU * load from metrics defaults to CPU * load from metrics defaults to CPU	2019-09-16 10:47:19 -04:00
Ananya Harsh Jha	c0f3b6b035	added set_epoch for distributed sampler, fix for #224 (#225 )	2019-09-16 10:21:00 -04:00
William Falcon	e339799a0a	Update README.md	2019-09-14 09:55:42 -04:00
William Falcon	50f5e4bec8	Update single_cpu_template.py	2019-09-14 02:23:49 -04:00
William Falcon	330a21ea91	Update README.md	2019-09-14 02:18:33 -04:00
William Falcon	f3221a5014	Update multi_node_cluster_auto_slurm.py	2019-09-14 02:14:08 -04:00
William Falcon	fe17d14ade	Update multi_node_cluster_auto_slurm.py	2019-09-13 17:05:49 -04:00
William Falcon	9576dd28b2	added load on CPU first (#221 ) * added load on CPU first * added load on CPU first * added load on CPU first * added load on CPU first * added load on CPU first * added load on CPU first * added load on CPU first * added load on CPU first * added load on CPU first * added load on CPU first * added load on CPU first * added load on CPU first * added load on CPU first * added load on CPU first * added load on CPU first * added load on CPU first * added load on CPU first * added load on CPU first * added load on CPU first * added load on CPU first * added load on CPU first * added load on CPU first * added load on CPU first * added load on CPU first * added load on CPU first * added load on CPU first * added load on CPU first * added print logs * added print logs * changed close order * changed close order	2019-09-11 07:52:36 -04:00
William Falcon	90353ac54e	changed examples scripts	2019-09-11 07:05:15 -04:00
William Falcon	cf7dbf6d7c	changed examples scripts	2019-09-11 07:03:31 -04:00
William Falcon	30b25c8146	Sai prasanna master (#219 ) * Fix incorrect warning for DistributedSampler. Check whether `dataloader.sampler` is an instance of DistributedSampler instead of checking the `dataloader`. * Update trainer.py * merged	2019-09-09 11:36:24 -04:00
William Falcon	ac0111c196	Update multi_node_cluster_auto_slurm.py	2019-09-09 10:55:47 -04:00
William Falcon	cbc619afa1	Update multi_node_own_slurm_script.py	2019-09-09 10:54:43 -04:00
William Falcon	3393086cb6	Update multi_node_cluster_auto_slurm.py	2019-09-09 10:53:47 -04:00
William Falcon	506d5da68b	enable single gpu per node (#218 ) * enable single gpu per node * enable single gpu per node * enable single gpu per node * enable single gpu per node * enable single gpu per node * enable single gpu per node	2019-09-09 07:37:20 -04:00
William Falcon	a6fe6f0917	Update README.md	2019-09-08 18:21:05 -04:00
William Falcon	8f289f9fa8	Update README.md	2019-09-08 18:19:00 -04:00
William Falcon	6c947f4e0d	Update README.md	2019-09-08 18:18:21 -04:00
William Falcon	396047ffa0	Updated distributed Demos (#215 ) * added simple cluster template * added simple cluster template * added simple cluster template * added simple cluster template * added simple cluster template * added simple cluster template * added simple cluster template * added simple cluster template * added simple cluster template * added simple cluster template * sets correct backend for possible combinations of gpu inputs * sets correct backend for possible combinations of gpu inputs * simple slurm example * simple slurm example * simple slurm example	2019-09-08 18:17:33 -04:00
William Falcon	83b756f77b	Update tox.ini	2019-09-08 15:46:30 -04:00
William Falcon	10d190e045	Simplified gpu api. No NVIDIA flag managing by lightning for cluster (#213 ) * added nvidia flag set * added nvidia flag set * added nvidia flag set * added nvidia flag set * added nvidia flag set * added nvidia flag set * added nvidia flag set * added nvidia flag set * added simple cluster template * sets correct backend for possible combinations of gpu inputs * sets correct backend for possible combinations of gpu inputs * sets correct backend for possible combinations of gpu inputs * sets correct backend for possible combinations of gpu inputs * sets correct backend for possible combinations of gpu inputs * sets correct backend for possible combinations of gpu inputs * sets correct backend for possible combinations of gpu inputs * sets correct backend for possible combinations of gpu inputs * sets correct backend for possible combinations of gpu inputs * sets correct backend for possible combinations of gpu inputs * sets correct backend for possible combinations of gpu inputs * sets correct backend for possible combinations of gpu inputs * sets correct backend for possible combinations of gpu inputs * sets correct backend for possible combinations of gpu inputs * sets correct backend for possible combinations of gpu inputs * sets correct backend for possible combinations of gpu inputs	2019-09-08 15:36:58 -04:00
William Falcon	b3434943c7	Update multi_node_cluster_template.py	2019-09-07 10:31:20 -04:00
Alok Singh	81df2259ef	Make print_nan_grads print grad (#208 ) This seems more useful for debugging.	2019-09-07 01:08:09 -04:00
williamFalcon	9f9d38673e	fixed demo	2019-09-06 16:26:46 -07:00
William Falcon	0c7fbc7178	Weights path (#211 ) * added docs. removed options. added weights_save option * removed old restore * cleaned up save path * cleaned up save path * flake8	2019-09-06 17:01:03 -04:00
William Falcon	3e74ea15d8	Fixes #120 (#210 )	2019-09-06 14:27:24 -04:00
William Falcon	7099f8dbfb	split trainer mixins (#209 ) * split trainer mixins * Update multi_node_cluster_template.py * Update single_cpu_template.py * Update single_gpu_node_16bit_template.py * Update single_gpu_node_ddp_template.py * Update single_gpu_node_dp_template.py * Update trainer_cpu_template.py * Update trainer_io.py * split trainer mixins * Update multi_node_cluster_template.py * deconflicted * deconflicted * deconflicted	2019-09-06 14:11:07 -04:00
William Falcon	60633eaa32	Moves hpc auto-resubmit to trainer from test-tube (#207 ) * added slurm signal handler * added restore weight functions * set slurm signal handling inside process * added resubmit docs * added resubmit docs * fixed missing param * Update trainer.py * fixed missing param * fixed missing param * debugging tests * debugging tests * debugging tests * debugging tests * debugging tests * debugging tests * debugging tests	2019-09-06 11:54:51 -04:00
Jirka Borovec	7ed928dfac	add PR template (#204 ) * add PR template * Update PULL_REQUEST_TEMPLATE.md	2019-09-06 10:12:06 -04:00
Nic Eggert	1733dba735	Pass outputs from all dataloaders to test_end and validation_end (#203 ) * Pass outputs from all dataloaders to test_end and validation_end * Update tests * Update docs * Update trainer.py * Update test_models.py	2019-09-06 07:37:25 -04:00
Jirka Borovec	447ed30716	extend pip install info (#194 ) * extend pip install info * Update README.md * Update README.md	2019-09-06 07:30:51 -04:00
William Falcon	7e0ac3149c	refactored init (#206 )	2019-09-06 00:29:38 -04:00
Thomas J Fan	bd50d9a2b4	DOC Adds reference to test-tube (#205 )	2019-09-05 21:13:49 -04:00
Jirka Borovec	5ef6fa5608	add osx to Travis (#202 ) * add CI macOS * add CI Windows * update CI * drop Win * update CI * update CI	2019-09-05 15:08:19 -04:00
Anton Konstantinov	34b824a9d3	Implement correct transfer to GPU for batches (#200 )	2019-09-05 07:13:06 -04:00
Thomas J Fan	62252cee58	STY Minor flake8 fix (#197 )	2019-09-04 17:46:56 -04:00
Max Horn	dac41030d4	Allow to deactivate GPU memory logging in Trainer (#190 ) * Allow to deactivate GPU memory logging in Trainer Adds the flag `log_gpu_memory` to Trainer to deactivate logging of GPU memory utilization. On some servers logging the GPU memory usage can significantly slow down training. * Update Logging.md * Update trainer.py	2019-09-04 10:43:46 -04:00
Verena Haunschmid	0872c32151	fix import in Tensorboard example (#193 )	2019-09-04 10:20:59 -04:00
Thomas J Fan	c766167773	DOC Minor import fix (#192 )	2019-09-04 06:17:54 -04:00
Nic Eggert	64688e1e15	Refactor test modules (#180 ) * Expectopatronum implement #89 (#182) * rename validate -> evaluate; implement test logic; allow multiple test_loaders * add test_step and test_end to LightningModule * add in_test_mode to pretraining to implement case 2 (test pretrained model) * fix code style issues * LightningTestModel: add optional second test set, implement test_step and test_end * implemented test for multiple test_dataloaders; fixed typo * add two test cases for #89 * add documentation for test_step, test_end; fix computation of loss in validation_step example * Update trainer.py * Update trainer.py * Update trainer.py * Update trainer.py * Update trainer.py * Update trainer.py * Added proper dp ddp routing calls for test mode * Update trainer.py * Update test_models.py * Update trainer.py * Update trainer.py * Update override_data_parallel.py * Update test_models.py * Update test_models.py * Update trainer.py * Update trainer.py * Update trainer.py * Update test_models.py * Update test_models.py * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * Update trainer.py * Update override_data_parallel.py * Update debug.py * Update lm_test_module.py * Update test_models.py * release v0.4.8 * Update README.md * add training loop docs * testing loop docs * testing loop docs * Convert __dataloader to _dataloader This will let inherited classes use it * Factor common test model setup into base class * Specialized test modules inherit from LightningTestModelBase * Fix __is_overriden so that it works with more complicated inheritance * Use mixins to add functionality to test models * Fix test with no val_dataloader * Remove unused imports * Get rid of wild card import * Update trainer.py * Update lm_test_module.py	2019-09-02 15:46:16 -04:00
William Falcon	c4ce347f3e	testing loop docs	2019-09-02 07:15:45 -04:00
William Falcon	8d6648e51d	Update README.md	2019-09-02 07:15:45 -04:00
William Falcon	9e6ce3b0d6	testing loop docs	2019-09-02 07:15:45 -04:00
William Falcon	a327596b79	add training loop docs	2019-09-02 07:15:45 -04:00
William Falcon	08a1ae8069	release v0.4.8	2019-09-02 07:15:45 -04:00
Verena Haunschmid	25d5b25792	Expectopatronum implement #89 (#182 ) * rename validate -> evaluate; implement test logic; allow multiple test_loaders * add test_step and test_end to LightningModule * add in_test_mode to pretraining to implement case 2 (test pretrained model) * fix code style issues * LightningTestModel: add optional second test set, implement test_step and test_end * implemented test for multiple test_dataloaders; fixed typo * add two test cases for #89 * add documentation for test_step, test_end; fix computation of loss in validation_step example * Update trainer.py * Update trainer.py * Update trainer.py * Update trainer.py * Update trainer.py * Update trainer.py * Added proper dp ddp routing calls for test mode * Update trainer.py * Update test_models.py * Update trainer.py * Update trainer.py * Update override_data_parallel.py * Update test_models.py * Update test_models.py * Update trainer.py * Update trainer.py * Update trainer.py * Update test_models.py * Update test_models.py * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * Update trainer.py * Update override_data_parallel.py * Update debug.py * Update lm_test_module.py * Update test_models.py	2019-09-02 07:15:27 -04:00
Stanislav	73cf47112e	Gradient accumulation callback (#150 ) * Gradient accumulation callback * little test case * typo * import fix * method name fix * fix epochs indexing from 1 * better code style * code style fix v2 :/ * change interface * fix Trainre new api in tests * trainer api bug fix * new raising error, new update method * extentions tests * a little better tests * typo fix * flack8 better * using scheduler for int and dict * typo * firs epoch bug fix * test update * empty dict exception * floats check * codestyle fix * grad counting test * someday, i will install normal linter * add more checks * Update test_models.py * Update test_models.py * Update test_models.py * Update test_models.py * Update test_models.py * Update test_models.py * Update test_models.py	2019-08-30 10:56:14 -04:00
Ir1dXD	c2247350bb	feat(val_sanity): enable skipping validation sanity (#176 ) * feat(val_sanity): enable skipping validation sanity when self.nb_sanity_val_steps is 0 * docs: elaborate on skipping	2019-08-28 06:41:31 -04:00
William Falcon	67c314272b	Update setup.py (#174 )	2019-08-27 18:07:33 -04:00
Ir1dXD	da4c1e3409	docs: add repo_name in the upright corner (#171 )	2019-08-27 16:46:18 -04:00
Jirka Borovec	cd89b4ef43	move GH docs (#168 )	2019-08-27 07:10:26 -04:00
Ir1dXD	6eb6daa278	enable highlight (#170 )	2019-08-27 07:09:46 -04:00