removed old template

2026-09-11 12:31:23 +08:00 · 2019-07-27 13:59:20 -04:00
parent 3ef6e6a572 7dd22f82c6
commit b58ca39621
45 changed files with 2921 additions and 473 deletions
@@ -9,6 +9,7 @@ model_weights/
 app/models/
 pip-wheel-metadata/
 test_tube_exp/
+tests/tests_tt_dir/

 # Byte-compiled / optimized / DLL files
 __pycache__/
@@ -119,4 +120,4 @@ ENV/
 .mypy_cache/

 # data
-mnist/
+mnist/
@@ -0,0 +1,19 @@
+# .readthedocs.yml
+# Read the Docs configuration file
+# See https://docs.readthedocs.io/en/stable/config-file/v2.html for details
+
+# Required
+version: 2
+
+# Build documentation with MkDocs
+mkdocs:
+  configuration: mkdocs.yml
+
+# Optionally build your docs in additional formats such as PDF and ePub
+formats: all
+
+# Optionally set the version of Python and requirements required to build your docs
+python:
+  version: 3.7
+  install:
+    - requirements: docs/doc_requirements.txt
@@ -0,0 +1,16 @@
+language: python
+python:
+  - "3.7"
+# command to install dependencies
+cache: pip
+install:
+  - pip install -e .
+  - pip install -r requirements.txt
+  - pip install -U numpy
+
+# keep build from timing out
+dist: xenial
+
+# command to run tests
+script:
+  - py.test # or py.test for Python versions 3.5 and below
@@ -9,31 +9,103 @@
 <p align="center">
  The Keras for ML researchers using PyTorch. More control. Less boilerplate.    
 </p>
+
 <p align="center">
  <a href="https://badge.fury.io/py/pytorch-lightning"><img src="https://badge.fury.io/py/pytorch-lightning.svg" alt="PyPI version" height="18"></a>
-<!--   <a href="https://travis-ci.org/williamFalcon/test-tube"><img src="https://travis-ci.org/williamFalcon/pytorch-lightning.svg?branch=master"></a> -->
+  <a href="https://pepy.tech/project/pytorch-lightning"><img src="https://pepy.tech/badge/pytorch-lightning" alt="PyPI version" height="18"></a>
+  <a href="https://github.com/williamFalcon/pytorch-lightning/tree/master/tests"><img src="https://github.com/williamFalcon/pytorch-lightning/blob/master/coverage.svg"></a>
+  <a href="https://travis-ci.org/williamFalcon/pytorch-lightning"><img src="https://travis-ci.org/williamFalcon/pytorch-lightning.svg?branch=master"></a>
+  <a href="https://williamfalcon.github.io/pytorch-lightning/"><img src="https://readthedocs.org/projects/pytorch-lightning/badge/?version=latest"></a>
  <a href="https://github.com/williamFalcon/pytorch-lightning/blob/master/COPYING"><img src="https://img.shields.io/badge/License-MIT-yellow.svg"></a>
 </p>   

 ```bash
-pip install pytorch-lightning    
+pip install pytorch-lightning  
 ```

 ## Docs   
 **[View the docs here](https://williamfalcon.github.io/pytorch-lightning/)**

 ## What is it?  
-Keras and fast.ai are too abstract for researchers. Lightning abstracts the full training loop but gives you control in the critical points.   
+Lightning defers training and validation loop logic to you. It guarantees correct, modern best practices for the core training logic.


 ## Why do I want to use lightning?
-Because you don't want to define a training loop, validation loop, gradient clipping, checkpointing, loading, 
-gpu training, etc... every time you start a project. Let lightning handle all of that for you! Just define your 
-data and what happens in the training, testing and validation loop and lightning will do the rest.
+When starting a new project the last thing you want to do is recode a training loop, model loading/saving, distributed training, when to validate, etc... You're likely to spend a long time ironing out all the bugs without even getting to the core of your research.
+
+With lightning, you guarantee those parts of your code work so you can focus on what the meat of the research: Data and training, validation loop logic. Don't worry about multiple gpus or speeding up your code, lightning will do that for you!
+
+## How do I do use it?   

 To use lightning do 2 things:  
-1. [Define a Trainer](https://github.com/williamFalcon/pytorch-lightning/blob/master/examples/new_project_templates/trainer_cpu_template.py).   
-2. [Define a LightningModel](https://github.com/williamFalcon/pytorch-lightning/blob/master/examples/new_project_templates/lightning_module_template.py).     
+1. [Define a LightningModel](https://williamfalcon.github.io/pytorch-lightning/LightningModule/RequiredTrainerInterface/)         
+```python
+import pytorch_lightning as ptl
+import torch
+from torch.nn import functional as F
+from torch.utils.data import DataLoader
+from torchvision.datasets import MNIST
+
+class CoolModel(ptl.LightningModule):
+
+    def __init(self):
+        super(CoolModel, self).__init__()
+        # not the best model...
+        self.l1 = torch.nn.Linear(28 * 28, 10)
+
+    def forward(self, x):
+        return torch.relu(self.l1(x))
+
+    def my_loss(self, y_hat, y):
+        return F.cross_entropy(y_hat, y)
+
+    def training_step(self, batch, batch_nb):
+        x, y = batch
+        y_hat = self.forward(x)
+        return {'tng_loss': self.my_loss(y_hat, y)}
+
+    def validation_step(self, batch, batch_nb):
+        x, y = batch
+        y_hat = self.forward(x)
+        return {'val_loss': self.my_loss(y_hat, y)}
+
+    def validation_end(self, outputs):
+        avg_loss = torch.stack([x for x in outputs['val_loss']]).mean()
+        return avg_loss
+
+    def configure_optimizers(self):
+        return [torch.optim.Adam(self.parameters(), lr=0.02)]
+
+    @ptl.data_loader
+    def tng_dataloader(self):
+        return DataLoader(MNIST('path/to/save', train=True), batch_size=32)
+
+    @ptl.data_loader
+    def val_dataloader(self):
+        return DataLoader(MNIST('path/to/save', train=False), batch_size=32)
+
+    @ptl.data_loader
+    def test_dataloader(self):
+        return DataLoader(MNIST('path/to/save', train=False), batch_size=32)
+```
+
+2. Fit with a [trainer](https://williamfalcon.github.io/pytorch-lightning/Trainer/)    
+```python
+from pytorch_lightning import Trainer
+from test_tube import Experiment
+
+model = CoolModel()
+
+# fit on 32 gpus across 4 nodes
+exp = Experiment(save_dir='some/dir')
+trainer = Trainer(experiment=exp, nb_gpu_nodes=4, gpus=[0,1,2,3,4,5,6,7])
+
+trainer.fit(model)
+
+# see all experiment metrics here
+# tensorboard --log_dir some/dir   
+```
+

 ## What does lightning control for me?
 Everything!    
@@ -116,8 +188,8 @@ def validation_end(self, outputs):
    return tqdm_dic
 ```
   
-## TensorboardX    
-Lightning is fully integrated with tensorboardX.   
+## Tensorboard    
+Lightning is fully integrated with tensorboard.   

 <p align="center">
  <a href="https://williamfalcon.github.io/pytorch-lightning/">
@@ -148,7 +220,7 @@ And run tensorboard from that dir
 tensorboard --logdir /some/path     
 ```    

-## Lightning automatically automates all of the following ([each is also configurable](https://williamfalcon.github.io/pytorch-lightning/Trainer/)):
+## Lightning automates all of the following ([each is also configurable](https://williamfalcon.github.io/pytorch-lightning/Trainer/)):

 ###### Checkpointing    

@@ -215,19 +287,22 @@ pip install pytorch-lightning

 # clone lightning for the demo
 git clone https://github.com/williamFalcon/pytorch-lightning.git
-cd examples/new_project_templates/
+cd pytorch_lightning/examples/new_project_templates/

-# run demo (on cpu)
-python trainer_gpu_cluster_template.py
+# all of the following demos use the SAME model to show no modification needs to be made to your code
+
+# train on cpu 
+python single_cpu_template.py
+
+# train on multiple-gpus 
+python single_gpu_node_template.py --gpus "0,1"
+
+# train on 32 gpus on a cluster (run on a SLURM managed cluster)
+python multi_node_cluster_template.py --nb_gpu_nodes 4 --gpus '0,1,2,3,4,5,6,7'
 ```

-Without changing the model AT ALL, you can run the model on a single gpu, over multiple gpus, or over multiple nodes.
+## Bleeding edge
+If you can't wait for the next release, install the most up to date code with:  
 ```bash
-# run a grid search on two gpus
-python fully_featured_trainer.py --gpus "0;1"
-
-# run single model on multiple gpus
-python fully_featured_trainer.py --gpus "0;1" --interactive
-```
-
-
+pip install git+https://github.com/williamFalcon/pytorch-lightning.git@master --upgrade
+```
@@ -0,0 +1,21 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<svg xmlns="http://www.w3.org/2000/svg" width="99" height="20">
+    <linearGradient id="b" x2="0" y2="100%">
+        <stop offset="0" stop-color="#bbb" stop-opacity=".1"/>
+        <stop offset="1" stop-opacity=".1"/>
+    </linearGradient>
+    <mask id="a">
+        <rect width="99" height="20" rx="3" fill="#fff"/>
+    </mask>
+    <g mask="url(#a)">
+        <path fill="#555" d="M0 0h63v20H0z"/>
+        <path fill="#4c1" d="M63 0h36v20H63z"/>
+        <path fill="url(#b)" d="M0 0h99v20H0z"/>
+    </g>
+    <g fill="#fff" text-anchor="middle" font-family="DejaVu Sans,Verdana,Geneva,sans-serif" font-size="11">
+        <text x="31.5" y="15" fill="#010101" fill-opacity=".3">coverage</text>
+        <text x="31.5" y="14">coverage</text>
+        <text x="80" y="15" fill="#010101" fill-opacity=".3">99%</text>
+        <text x="80" y="14">99%</text>
+    </g>
+</svg>
@@ -3,7 +3,7 @@

 A lightning module is a strict superclass of nn.Module, it provides a standard interface for the trainer to interact with the model.

-The easiest thing to do is copy [this template](../../examples/new_project_templates/lightning_module_template.py) and modify accordingly. 
+The easiest thing to do is copy [this template](../../pytorch_lightning/examples/new_project_templates/lightning_module_template.py) and modify accordingly. 

 Otherwise, to Define a Lightning Module, implement the following methods:

@@ -14,8 +14,6 @@ Otherwise, to Define a Lightning Module, implement the following methods:
 - [validation_end](RequiredTrainerInterface.md#validation_end)
    
 - [configure_optimizers](RequiredTrainerInterface.md#configure_optimizers)
- [get_save_dict](RequiredTrainerInterface.md#get_save_dict)
- [load_model_specific](RequiredTrainerInterface.md#load_model_specific)

 - [tng_dataloader](RequiredTrainerInterface.md#tng_dataloader)
 - [tng_dataloader](RequiredTrainerInterface.md#tng_dataloader)
@@ -23,9 +21,63 @@ Otherwise, to Define a Lightning Module, implement the following methods:

 **Optional**:   

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

+---
+**Minimal example**
+```python
+import pytorch_lightning as ptl
+import torch
+from torch.nn import functional as F
+from torch.utils.data import DataLoader
+from torchvision.datasets import MNIST
+
+class CoolModel(ptl.LightningModule):
+
+    def __init(self):
+        super(CoolModel, self).__init__()
+        # not the best model...
+        self.l1 = torch.nn.Linear(28 * 28, 10)
+
+    def forward(self, x):
+        return torch.relu(self.l1(x))
+
+    def my_loss(self, y_hat, y):
+        return F.cross_entropy(y_hat, y)
+
+    def training_step(self, batch, batch_nb):
+        x, y = batch
+        y_hat = self.forward(x)
+        return {'tng_loss': self.my_loss(y_hat, y)}
+
+    def validation_step(self, batch, batch_nb):
+        x, y = batch
+        y_hat = self.forward(x)
+        return {'val_loss': self.my_loss(y_hat, y)}
+
+    def validation_end(self, outputs):
+        avg_loss = torch.stack([x for x in outputs['val_loss']]).mean()
+        return avg_loss
+
+    def configure_optimizers(self):
+        return [torch.optim.Adam(self.parameters(), lr=0.02)]
+
+    @ptl.data_loader
+    def tng_dataloader(self):
+        return DataLoader(MNIST('path/to/save', train=True), batch_size=32)
+
+    @ptl.data_loader
+    def val_dataloader(self):
+        return DataLoader(MNIST('path/to/save', train=False), batch_size=32)
+
+    @ptl.data_loader
+    def test_dataloader(self):
+        return DataLoader(MNIST('path/to/save', train=False), batch_size=32)
+```
+
 ---

 ### training_step
@@ -193,34 +245,35 @@ def configure_optimizers(self):
 ```

 --- 
-### get_save_dict 
+### on_save_checkpoint 

 ``` {.python}
-def get_save_dict(self)
+def on_save_checkpoint(self, checkpoint)
 ```
-Called by lightning to checkpoint your model. Lightning saves current epoch, current batch nb, etc...
-All you have to return is what specifically about your lightning model you want to checkpoint.
+Called by lightning to checkpoint your model. Lightning saves the training state (current epoch, global_step, etc)
+and also saves the model state_dict. If you want to save anything else, use this method to add your own
+key-value pair.

 ##### Return
-Dictionary - No required keys. Most of the time as described in this example.   
+Nothing

 **Example**

 ``` {.python}
-def get_save_dict(self):
-    # 99% of use cases this is all you need to return
-    checkpoint = {'state_dict': self.state_dict()}
-    return checkpoint
+def on_save_checkpoint(self, checkpoint):
+    # 99% of use cases you don't need to implement this method 
+    checkpoint['something_cool_i_want_to_save'] = my_cool_pickable_object
 ```

 --- 
-### load_model_specific 
+### on_load_checkpoint 

 ``` {.python}
-def load_model_specific(self, checkpoint)
+def on_load_checkpoint(self, checkpoint)
 ```
-Called by lightning to restore your model. This is your chance to restore your model using the keys you added in get_save_dict.
-Lightning will automatically restore current epoch, batch nb, etc. 
+Called by lightning to restore your model. Lighting auto-restores global step, epoch, etc...
+It also restores the model state_dict.
+If you saved something with **on_save_checkpoint** this is your chance to restore this.

 ##### Return
 Nothing  
@@ -228,19 +281,19 @@ Nothing
 **Example**

 ``` {.python}
-def load_model_specific(self, checkpoint):
-    # you defined 'state_dict' in get_save_dict()
-    self.load_state_dict(checkpoint['state_dict'])
+def on_load_checkpoint(self, checkpoint):
+    # 99% of the time you don't need to implement this method
+    self.something_cool_i_want_to_save = checkpoint['something_cool_i_want_to_save']
 ```

 --- 
 ### tng_dataloader 

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

 ##### Return
 Pytorch DataLoader
@@ -248,32 +301,26 @@ Pytorch DataLoader
 **Example**

 ``` {.python}
-@property
+@ptl.data_loader
 def tng_dataloader(self):
-    if self._tng_dataloader is None:
-        try:
-            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(
-                dataset=dataset,
-                batch_size=self.hparams.batch_size,
-                shuffle=True
-            )
-            self._tng_dataloader = loader
-        except Exception as e:
-            raise e
-            
-    return self._tng_dataloader
+    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(
+        dataset=dataset,
+        batch_size=self.hparams.batch_size,
+        shuffle=True
+    )
+    return loader
 ```

 --- 
 ### val_dataloader 

 ``` {.python}
-@property
+@ptl.data_loader
 def tng_dataloader(self)
 ```
-Called by lightning during validation loop. Define it as a property.
+Called by lightning during validation loop. Make sure to use the @ptl.data_loader decorator, this ensures not calling this function until the data are needed.

 ##### Return
 Pytorch DataLoader
@@ -281,32 +328,27 @@ Pytorch DataLoader
 **Example**

 ``` {.python}
-@property
+@ptl.data_loader
 def val_dataloader(self):
-    if self._val_dataloader is None:
-        try:
-            transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,), (1.0,))])
-            dataset = MNIST(root='/path/to/mnist/', train=False, transform=transform, download=True)
-            loader = torch.utils.data.DataLoader(
-                dataset=dataset,
-                batch_size=self.hparams.batch_size,
-                shuffle=True
-            )
-            self._val_dataloader = loader
-        except Exception as e:
-            raise e
-            
-    return self._val_dataloader
+    transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,), (1.0,))])
+    dataset = MNIST(root='/path/to/mnist/', train=False, transform=transform, download=True)
+    loader = torch.utils.data.DataLoader(
+        dataset=dataset,
+        batch_size=self.hparams.batch_size,
+        shuffle=True
+    )
+    
+    return loader
 ```

 --- 
 ### test_dataloader 

 ``` {.python}
-@property
+@ptl.data_loader
 def test_dataloader(self)
 ```
-Called by lightning during test loop. Define it as a property.
+Called by lightning during test loop. Make sure to use the @ptl.data_loader decorator, this ensures not calling this function until the data are needed.

 ##### Return
 Pytorch DataLoader
@@ -314,22 +356,17 @@ Pytorch DataLoader
 **Example**

 ``` {.python}
-@property
+@ptl.data_loader
 def test_dataloader(self):
-    if self._test_dataloader is None:
-        try:
-            transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,), (1.0,))])
-            dataset = MNIST(root='/path/to/mnist/', train=False, transform=transform, download=True)
-            loader = torch.utils.data.DataLoader(
-                dataset=dataset,
-                batch_size=self.hparams.batch_size,
-                shuffle=True
-            )
-            self._test_dataloader = loader
-        except Exception as e:
-            raise e
-            
-    return self._test_dataloader
+    transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,), (1.0,))])
+    dataset = MNIST(root='/path/to/mnist/', train=False, transform=transform, download=True)
+    loader = torch.utils.data.DataLoader(
+        dataset=dataset,
+        batch_size=self.hparams.batch_size,
+        shuffle=True
+    )
+    
+    return loader
 ```

 --- 
@@ -21,7 +21,8 @@ pretrained_model = MyLightningModule.load_from_metrics(
    map_location=None
 )
    
-# predict 
+# predict
+pretrained_model.eval()
 pretrained_model.freeze()
 y_hat = pretrained_model(x)
 ```
@@ -3,6 +3,26 @@ Lightning makes multi-gpu training and 16 bit training trivial.
 *Note:*   
 None of the flags below require changing anything about your lightningModel definition. 

+---
+#### Choosing a backend  
+Lightning supports two backends. DataParallel and DistributedDataParallel. Both can be used for single-node multi-GPU training.
+For multi-node training you must use DistributedDataParallel.   
+
+You can toggle between each mode by setting this flag.
+``` {.python}
+# DEFAULT uses DataParallel
+trainer = Trainer(distributed_backend='dp')
+
+# change to distributed data parallel
+trainer = Trainer(distributed_backend='ddp')
+```
+
+If you request multiple nodes, the back-end will auto-switch to ddp.
+We recommend you use DistributedDataparallel even for single-node multi-GPU training. It is MUCH faster than DP but *may*
+have configuration issues depending on your cluster.
+
+For a deeper understanding of what lightning is doing, feel free to read [this guide](https://medium.com/@_willfalcon/9-tips-for-training-lightning-fast-neural-networks-in-pytorch-8e63a502f565).   
+
 ---
 #### 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.    
@@ -37,16 +57,69 @@ 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
 # set these flags
-os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
-os.environ["CUDA_VISIBLE_DEVICES"] = "0,1,2,3,4,5,6,7"
+# lightning sets these flags for you automatically
+# no need to set yourself
+# os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
+# os.environ["CUDA_VISIBLE_DEVICES"] = "0,1,2,3,4,5,6,7"

-# DEFAULT
-trainer = Trainer(gpus=[0,1,2,3,4,5,6,7])
+
+# to use DataParallel (default)
+trainer = Trainer(gpus=[0,1,2,3,4,5,6,7], distributed_backend='dp')
+
+# RECOMMENDED use DistributedDataParallel
+trainer = Trainer(gpus=[0,1,2,3,4,5,6,7], distributed_backend='ddp')
 ```

 ---
 #### Multi-node
-COMING SOON.
+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)
+```
+
+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.
+
+```python
+cluster = SlurmCluster(
+    hyperparam_optimizer=test_tube.HyperOptArgumentParser(),
+    log_path='/some/path/to/save',
+)
+
+# OPTIONAL FLAGS WHICH MAY BE CLUSTER DEPENDENT
+# which interface your nodes use for communication
+cluster.add_command('export NCCL_SOCKET_IFNAME=^docker0,lo')
+
+# see output of the NCCL connection process
+# NCCL is how the nodes talk to each other
+cluster.add_command('export NCCL_DEBUG=INFO')
+
+# setting a master port here is a good idea.
+cluster.add_command(f'export MASTER_PORT={PORT}')
+
+# good to load the latest NCCL version
+cluster.load_modules(['NCCL/2.4.7-1-cuda.10.0'])
+
+# configure cluster
+cluster.per_experiment_nb_nodes = 12 
+cluster.per_experiment_nb_gpus = 8
+
+cluster.add_slurm_cmd(cmd='ntasks-per-node', value=8, comment='1 task per gpu')
+```
+
+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).   
+
+```python
+# ie: this:
+dataset = myDataset()
+dataloader = Dataloader(dataset)
+
+# becomes:
+dataset = myDataset()
+dist_sampler = torch.utils.data.distributed.DistributedSampler(dataset)
+dataloader = Dataloader(dataset, sampler=dist_sampler)
+```

 ---
 #### Self-balancing architecture
@@ -19,7 +19,7 @@ Cut the learning rate by 10 at every epoch listed in this list.
 trainer = Trainer(lr_scheduler_milestones=None)

 # cut LR by 10 at 100, 200, and 300 epochs 
-trainer = Trainer(lr_scheduler_milestones=[100, 200, 300])
+trainer = Trainer(lr_scheduler_milestones='100, 200, 300')
 ```

 ---
@@ -0,0 +1 @@
+mkdocs-material==4.4.0
@@ -1,10 +1,11 @@
 ###### New project Quick Start    
 To start a new project define these two files.   

-1. [Define a LightningModule](/LightningModule/RequiredTrainerInterface/#template-model-definition)  
-2. Pick a trainer      
-    - [Basic CPU Trainer](https://github.com/williamFalcon/pytorch-lightning/blob/master/examples/new_project_templates/trainer_cpu_template.py) 
-    - [GPU cluster Trainer](https://github.com/williamFalcon/pytorch-lightning/blob/master/examples/new_project_templates/trainer_gpu_cluster_template.py)
+1. [Define a LightningModule](/pytorch-lightning/LightningModule/RequiredTrainerInterface/) 
+2. [Define a trainer](https://williamfalcon.github.io/pytorch-lightning/Trainer/)
+    - [Basic CPU Trainer Template](https://github.com/williamFalcon/pytorch-lightning/blob/master/pytorch_lightning/examples/new_project_templates/single_cpu_template.py) 
+    - [Multi-GPU Trainer Template](https://github.com/williamFalcon/pytorch-lightning/blob/master/pytorch_lightning/examples/new_project_templates/single_gpu_node_template.py)
+    - [GPU cluster Trainer Template](https://github.com/williamFalcon/pytorch-lightning/blob/master/pytorch_lightning/examples/new_project_templates/multi_node_cluster_template.py)

 ###### Docs shortcuts
 - [LightningModule](LightningModule/RequiredTrainerInterface/)  
@@ -1 +0,0 @@
-from .lightning_module_template import LightningTemplateModel
@@ -1,2 +1,3 @@
 from .models import Trainer
-from .root_module.root_module import LightningModule
+from .root_module.root_module import LightningModule
+from .root_module.decorators import data_loader
@@ -0,0 +1 @@
+from .new_project_templates.lightning_module_template import LightningTemplateModel
@@ -10,6 +10,7 @@ from torch import optim
 from torch.utils.data import DataLoader
 from torch.utils.data.distributed import DistributedSampler

+import pytorch_lightning as ptl
 from pytorch_lightning.root_module.root_module import LightningModule


@@ -24,10 +25,14 @@ class LightningTemplateModel(LightningModule):
        :param hparams:
        """
        # init superclass
-        super(LightningTemplateModel, self).__init__(hparams)
+        super(LightningTemplateModel, 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)
+
        # build model
        self.__build_model()

@@ -78,15 +83,21 @@ class LightningTemplateModel(LightningModule):
        # 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)
+
        output = OrderedDict({
-            'loss': loss_val,
-            'tqdm_metrics': {}
+            'loss': loss_val
        })
+
+        # can also return just a scalar instead of a dict (return loss_val)
        return output

    def validation_step(self, data_batch, batch_i):
@@ -104,11 +115,22 @@ class LightningTemplateModel(LightningModule):
        # 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)

        output = OrderedDict({
            'val_loss': loss_val,
-            'val_acc': torch.tensor(val_acc),
+            'val_acc': val_acc,
        })
+
+        # can also return just a scalar instead of a dict (return loss_val)
        return output

    def validation_end(self, outputs):
@@ -117,6 +139,10 @@ class LightningTemplateModel(LightningModule):
        :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:
@@ -128,20 +154,6 @@ class LightningTemplateModel(LightningModule):
        tqdm_dic = {'val_loss': val_loss_mean.item(), 'val_acc': val_acc_mean.item()}
        return tqdm_dic

-    def update_tng_log_metrics(self, logs):
-        return logs
-
-    # ---------------------
-    # MODEL SAVING
-    # ---------------------
-    def get_save_dict(self):
-        checkpoint = {'state_dict': self.state_dict()}
-        return checkpoint
-
-    def load_model_specific(self, checkpoint):
-        self.load_state_dict(checkpoint['state_dict'])
-        pass
-
    # ---------------------
    # TRAINING SETUP
    # ---------------------
@@ -179,38 +191,23 @@ class LightningTemplateModel(LightningModule):

        return loader

-    @property
+    @ptl.data_loader
    def tng_dataloader(self):
-        if self._tng_dataloader is None:
-            try:
-                self._tng_dataloader = self.__dataloader(train=True)
-            except Exception as e:
-                print(e)
-                raise e
-        return self._tng_dataloader
+        print('tng data loader called')
+        return self.__dataloader(train=True)

-    @property
+    @ptl.data_loader
    def val_dataloader(self):
-        if self._val_dataloader is None:
-            try:
-                self._val_dataloader = self.__dataloader(train=False)
-            except Exception as e:
-                print(e)
-                raise e
-        return self._val_dataloader
+        print('val data loader called')
+        return self.__dataloader(train=False)

-    @property
+    @ptl.data_loader
    def test_dataloader(self):
-        if self._test_dataloader is None:
-            try:
-                self._test_dataloader = self.__dataloader(train=False)
-            except Exception as e:
-                print(e)
-                raise e
-        return self._test_dataloader
+        print('test data loader called')
+        return self.__dataloader(train=False)

    @staticmethod
-    def add_model_specific_args(parent_parser, root_dir):
+    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:
@@ -0,0 +1,112 @@
+"""
+Runs a model on a single node across N-gpus.
+"""
+import os
+import sys
+import numpy as np
+from time import sleep
+import torch
+
+from test_tube import HyperOptArgumentParser, Experiment, SlurmCluster
+from pytorch_lightning.models.trainer import Trainer
+from pytorch_lightning.utils.arg_parse import add_default_args
+
+from pytorch_lightning.callbacks import EarlyStopping, ModelCheckpoint
+
+SEED = 2334
+torch.manual_seed(SEED)
+np.random.seed(SEED)
+
+from lightning_module_template import LightningTemplateModel
+
+
+def main(hparams):
+    """
+    Main training routine specific for this project
+    :param hparams:
+    :return:
+    """
+    # ------------------------
+    # 1 INIT LIGHTNING MODEL
+    # ------------------------
+    print('loading model...')
+    model = LightningTemplateModel(hparams)
+    print('model built')
+
+    # ------------------------
+    # 2 INIT TEST TUBE EXP
+    # ------------------------
+
+    # init experiment
+    exp = Experiment(
+        name=hyperparams.experiment_name,
+        save_dir=hyperparams.test_tube_save_path,
+        autosave=False,
+        description='test demo'
+    )
+
+    exp.argparse(hparams)
+    exp.save()
+
+    # ------------------------
+    # 3 DEFINE CALLBACKS
+    # ------------------------
+    model_save_path = '{}/{}/{}'.format(hparams.model_save_path, exp.name, exp.version)
+    early_stop = EarlyStopping(
+        monitor='val_acc',
+        patience=3,
+        verbose=True,
+        mode='max'
+    )
+
+    checkpoint = ModelCheckpoint(
+        filepath=model_save_path,
+        save_best_only=True,
+        verbose=True,
+        monitor='val_loss',
+        mode='min'
+    )
+
+    # ------------------------
+    # 4 INIT TRAINER
+    # ------------------------
+    trainer = Trainer(
+        experiment=exp,
+        checkpoint_callback=checkpoint,
+        early_stop_callback=early_stop,
+        gpus=hparams.gpus,
+    )
+
+    # ------------------------
+    # 5 START TRAINING
+    # ------------------------
+    trainer.fit(model)
+
+
+if __name__ == '__main__':
+
+    # dirs
+    root_dir = os.path.dirname(os.path.realpath(__file__))
+    demo_log_dir = os.path.join(root_dir, 'pt_lightning_demo_logs')
+    checkpoint_dir = os.path.join(demo_log_dir, 'model_weights')
+    test_tube_dir = os.path.join(demo_log_dir, 'test_tube_data')
+
+    # although we user hyperOptParser, we are using it only as argparse right now
+    parent_parser = HyperOptArgumentParser(strategy='grid_search', add_help=False)
+
+    # gpu args
+    parent_parser.add_argument('--gpus', type=str, default='-1', help='how many gpus to use in the node. -1 uses all the gpus on the node')
+    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, help='where to save model')
+    parent_parser.add_argument('--experiment_name', type=str, default='pt_lightning_exp_a', help='test tube exp name')
+
+    # allow model to overwrite or extend args
+    parser = LightningTemplateModel.add_model_specific_args(parent_parser, root_dir)
+    hyperparams = parser.parse_args()
+
+    # ---------------------
+    # RUN TRAINING
+    # ---------------------
+    # run on HPC cluster
+    print(f'RUNNING INTERACTIVE MODE ON GPUS. gpu ids: {hyperparams.gpus}')
+    main(hyperparams)
@@ -0,0 +1,112 @@
+"""
+Runs a model on a single node across N-gpus.
+"""
+import os
+import sys
+import numpy as np
+from time import sleep
+import torch
+
+from test_tube import HyperOptArgumentParser, Experiment, SlurmCluster
+from pytorch_lightning.models.trainer import Trainer
+from pytorch_lightning.utils.arg_parse import add_default_args
+
+from pytorch_lightning.callbacks import EarlyStopping, ModelCheckpoint
+
+SEED = 2334
+torch.manual_seed(SEED)
+np.random.seed(SEED)
+
+from lightning_module_template import LightningTemplateModel
+
+
+def main(hparams):
+    """
+    Main training routine specific for this project
+    :param hparams:
+    :return:
+    """
+    # ------------------------
+    # 1 INIT LIGHTNING MODEL
+    # ------------------------
+    print('loading model...')
+    model = LightningTemplateModel(hparams)
+    print('model built')
+
+    # ------------------------
+    # 2 INIT TEST TUBE EXP
+    # ------------------------
+
+    # init experiment
+    exp = Experiment(
+        name=hyperparams.experiment_name,
+        save_dir=hyperparams.test_tube_save_path,
+        autosave=False,
+        description='test demo'
+    )
+
+    exp.argparse(hparams)
+    exp.save()
+
+    # ------------------------
+    # 3 DEFINE CALLBACKS
+    # ------------------------
+    model_save_path = '{}/{}/{}'.format(hparams.model_save_path, exp.name, exp.version)
+    early_stop = EarlyStopping(
+        monitor='val_acc',
+        patience=3,
+        verbose=True,
+        mode='max'
+    )
+
+    checkpoint = ModelCheckpoint(
+        filepath=model_save_path,
+        save_best_only=True,
+        verbose=True,
+        monitor='val_loss',
+        mode='min'
+    )
+
+    # ------------------------
+    # 4 INIT TRAINER
+    # ------------------------
+    trainer = Trainer(
+        experiment=exp,
+        checkpoint_callback=checkpoint,
+        early_stop_callback=early_stop,
+        gpus=hparams.gpus,
+    )
+
+    # ------------------------
+    # 5 START TRAINING
+    # ------------------------
+    trainer.fit(model)
+
+
+if __name__ == '__main__':
+
+    # dirs
+    root_dir = os.path.dirname(os.path.realpath(__file__))
+    demo_log_dir = os.path.join(root_dir, 'pt_lightning_demo_logs')
+    checkpoint_dir = os.path.join(demo_log_dir, 'model_weights')
+    test_tube_dir = os.path.join(demo_log_dir, 'test_tube_data')
+
+    # although we user hyperOptParser, we are using it only as argparse right now
+    parent_parser = HyperOptArgumentParser(strategy='grid_search', add_help=False)
+
+    # gpu args
+    parent_parser.add_argument('--gpus', type=str, default='0', help='how many gpus to use in the node. -1 uses all the gpus on the node')
+    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, help='where to save model')
+    parent_parser.add_argument('--experiment_name', type=str, default='pt_lightning_exp_a', help='test tube exp name')
+
+    # allow model to overwrite or extend args
+    parser = LightningTemplateModel.add_model_specific_args(parent_parser, root_dir)
+    hyperparams = parser.parse_args()
+
+    # ---------------------
+    # RUN TRAINING
+    # ---------------------
+    # run on HPC cluster
+    print(f'RUNNING INTERACTIVE MODE ON GPUS. gpu ids: {hyperparams.gpus}')
+    main(hyperparams)
@@ -0,0 +1,891 @@
+"""
+The trainer handles all the logic for running a val loop, training loop, distributing, etc...
+"""
+import subprocess
+import traceback
+import warnings
+import os
+import pdb
+import re
+
+import torch
+from torch.utils.data.distributed import DistributedSampler
+from torch.optim.lr_scheduler import MultiStepLR
+import torch.multiprocessing as mp
+import torch.distributed as dist
+import numpy as np
+import tqdm
+
+from pytorch_lightning.root_module.memory import get_gpu_memory_map
+from pytorch_lightning.root_module.model_saving import TrainerIO
+from pytorch_lightning.pt_overrides.override_data_parallel import LightningDistributedDataParallel, LightningDataParallel
+from pytorch_lightning.utils.debugging import MisconfigurationException
+
+try:
+    from apex import amp
+    APEX_AVAILABLE = True
+except Exception:
+    APEX_AVAILABLE = False
+
+
+def reduce_distributed_output(output, nb_gpus):
+    if nb_gpus <= 1:
+        return output
+
+    # when using DP, we get one output per gpu
+    # average outputs and return
+    if type(output) is torch.Tensor:
+        return output.mean()
+
+    for k, v in output.items():
+        # recurse on nested dics
+        if isinstance(output[k], dict):
+            output[k] = reduce_distributed_output(output[k], nb_gpus)
+
+        # reduce only metrics that have the same nb of gpus
+        elif output[k].size(0) == nb_gpus:
+            reduced = torch.mean(output[k])
+            output[k] = reduced
+    return output
+
+
+class Trainer(TrainerIO):
+
+    def __init__(self,
+                 experiment,
+                 early_stop_callback=None,
+                 checkpoint_callback=None,
+                 gradient_clip=0,
+                 cluster=None,
+                 process_position=0,
+                 current_gpu_name=0,
+                 nb_gpu_nodes=1,
+                 gpus=None,
+                 progress_bar=True,
+                 overfit_pct=0.0,
+                 track_grad_norm=-1,
+                 check_val_every_n_epoch=1,
+                 fast_dev_run=False,
+                 accumulate_grad_batches=1,
+                 max_nb_epochs=1000, min_nb_epochs=1,
+                 train_percent_check=1.0, val_percent_check=1.0, test_percent_check=1.0,
+                 val_check_interval=0.95,
+                 log_save_interval=100, add_log_row_interval=10,
+                 lr_scheduler_milestones=None,
+                 distributed_backend='dp',
+                 use_amp=False,
+                 print_nan_grads=False,
+                 print_weights_summary=True,
+                 amp_level='O2',
+                 nb_sanity_val_steps=5):
+
+        """
+
+        :param experiment: Test-tube experiment
+        :param early_stop_callback: from pytorch_lightning import EarlyStopping
+        :param checkpoint_callback: from pytorch_lightning import Checkpoint
+        :param gradient_clip:
+        :param cluster:
+        :param process_position:
+        :param current_gpu_name:
+        :param nb_gpu_nodes:
+        :param gpus:
+        :param progress_bar:
+        :param overfit_pct:
+        :param track_grad_norm:
+        :param check_val_every_n_epoch:
+        :param fast_dev_run:
+        :param accumulate_grad_batches:
+        :param max_nb_epochs:
+        :param min_nb_epochs:
+        :param train_percent_check:
+        :param val_percent_check:
+        :param test_percent_check:
+        :param val_check_interval:
+        :param log_save_interval:
+        :param add_log_row_interval:
+        :param lr_scheduler_milestones:
+        :param distributed_backend: 'np' to use DistributedParallel, 'ddp' to use DistributedDataParallel
+        :param use_amp:
+        :param print_nan_grads:
+        :param print_weights_summary:
+        :param amp_level:
+        :param nb_sanity_val_steps:
+        """
+
+        # Transfer params
+
+        self.nb_gpu_nodes = nb_gpu_nodes
+        self.gradient_clip = gradient_clip
+        self.check_val_every_n_epoch = check_val_every_n_epoch
+        self.enable_early_stop = early_stop_callback is not None
+        self.track_grad_norm = track_grad_norm
+        self.fast_dev_run = fast_dev_run
+        self.on_gpu = gpus is not None and torch.cuda.is_available()
+        self.progress_bar = progress_bar
+        self.experiment = experiment
+        self.exp_save_path = experiment.get_data_path(experiment.name, experiment.version)
+        self.cluster = cluster
+        self.process_position = process_position
+        self.current_gpu_name = current_gpu_name
+        self.print_weights_summary = print_weights_summary
+        self.checkpoint_callback = checkpoint_callback
+
+        if self.checkpoint_callback is not None:
+            self.checkpoint_callback.save_function = self.save_checkpoint
+
+        self.early_stop = early_stop_callback
+        self.model = None
+        self.max_nb_epochs = max_nb_epochs
+        self.accumulate_grad_batches = accumulate_grad_batches
+        self.early_stop_callback = early_stop_callback
+        self.min_nb_epochs = min_nb_epochs
+        self.nb_sanity_val_steps = nb_sanity_val_steps
+        self.lr_scheduler_milestones = [] if lr_scheduler_milestones is None else [int(x.strip()) for x in lr_scheduler_milestones.split(',')]
+        self.lr_schedulers = []
+        self.amp_level = amp_level
+        self.print_nan_grads = print_nan_grads
+        self.data_parallel_device_ids = None
+        self.world_size = 1
+        self.node_rank = 0
+        self.use_ddp = False
+        self.use_dp = False
+
+        # training bookeeping
+        self.total_batch_nb = 0
+        self.running_loss = []
+        self.avg_loss = 0
+        self.batch_nb = 0
+        self.tqdm_metrics = {}
+        self.nb_val_batches = None
+        self.nb_tng_batches = None
+        self.nb_test_batches = None
+
+        # gpus come in as a string.
+        # if gpus = -1 then use all available devices
+        # otherwise, split the string using commas
+        if gpus is not None:
+            if type(gpus) is list:
+                self.data_parallel_device_ids = gpus
+            elif type(gpus) is str:
+                if gpus == '-1':
+                    self.data_parallel_device_ids = list(range(0, torch.cuda.device_count()))
+                else:
+                    self.data_parallel_device_ids = [int(x.strip()) for x in gpus.split(',')]
+            else:
+                raise Exception('gpus has to be a string or list of ids')
+
+            # set the correct cuda visible devices (using pci order)
+            os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
+            os.environ["CUDA_VISIBLE_DEVICES"] = ','.join([str(x) for x in self.data_parallel_device_ids])
+            print(f'VISIBLE GPUS: {os.environ["CUDA_VISIBLE_DEVICES"]}')
+
+        # make DP and DDP mutually exclusive
+        # single GPU will also use DP with devices=[0]
+        have_gpus = self.data_parallel_device_ids is not None and len(self.data_parallel_device_ids) > 0
+        if have_gpus:
+            self.use_dp = distributed_backend == 'dp'
+            self.use_ddp = distributed_backend == 'ddp'
+
+            # use ddp automatically if nb_gpu_nodes > 1
+            if nb_gpu_nodes > 1 and self.use_dp:  # pragma: no cover
+                self.use_ddp = True
+                self.use_dp = False
+                w = 'DataParallel does not support nb_gpu_nodes > 1. ' \
+                    'Switching to DistributedDataParallel for you. ' \
+                    'To silence this warning set distributed_backend=ddp'
+                warnings.warn(w)
+
+        # extract SLURM flag vars
+        # whenever we have the correct number of tasks, we let slurm manage processes
+        # otherwise we launch the required number of processes
+        if self.use_ddp:
+            self.nb_requested_gpus = len(self.data_parallel_device_ids) * self.nb_gpu_nodes
+            self.nb_slurm_tasks = 0
+            try:
+                self.nb_slurm_tasks = int(os.environ['SLURM_NTASKS'])
+                self.is_slurm_managing_tasks = self.nb_slurm_tasks == self.nb_requested_gpus
+            except Exception as e:
+                # likely not on slurm, so set the slurm managed flag to false
+                self.is_slurm_managing_tasks = False
+
+        # process info
+        self.proc_rank = 0
+
+        # training state
+        self.optimizers = None
+        self.prog_bar = None
+        self.global_step = 0
+        self.current_epoch = 0
+        self.total_batches = 0
+
+        # logging
+        self.log_save_interval = log_save_interval
+        self.val_check_interval = val_check_interval
+        self.add_log_row_interval = add_log_row_interval
+
+        # dataloaders
+        self.tng_dataloader = None
+        self.test_dataloader = None
+        self.val_dataloader = None
+
+        # how much of the data to use
+        self.__determine_data_use_amount(train_percent_check, val_percent_check, test_percent_check, overfit_pct)
+        print('gpu available: {}, used: {}'.format(torch.cuda.is_available(), self.on_gpu))
+
+        # 16 bit mixed precision training using apex
+        self.use_amp = use_amp and APEX_AVAILABLE
+        if self.use_amp:
+            print('using 16bit precision')
+
+        if use_amp and not APEX_AVAILABLE:  # pragma: no cover
+            msg = '''
+            You set use_amp=True but do not have apex installed.
+            Install apex first using this guide and rerun with use_amp=True: 
+            https://github.com/NVIDIA/apex#linux
+            
+            this run will NOT use 16 bit precision
+            '''
+            raise ModuleNotFoundError(msg)
+
+    @property
+    def data_parallel(self):
+        return self.use_dp or self.use_ddp
+
+    def __determine_data_use_amount(self, train_percent_check, val_percent_check, test_percent_check, overfit_pct):
+        """
+        Use less data for debugging purposes
+        """
+        self.train_percent_check = train_percent_check
+        self.val_percent_check = val_percent_check
+        self.test_percent_check = test_percent_check
+        if overfit_pct > 0:
+            self.train_percent_check = overfit_pct
+            self.val_percent_check = overfit_pct
+            self.test_percent_check = overfit_pct
+
+    def __get_model(self):
+        return self.model.module if self.data_parallel else self.model
+
+    def __is_function_implemented(self, f_name):
+        model = self.__get_model()
+        f_op = getattr(model, f_name, None)
+        return callable(f_op)
+
+    @property
+    def __tng_tqdm_dic(self):
+        # ForkedPdb().set_trace()
+        tqdm_dic = {
+            'tng_loss': '{0:.3f}'.format(self.avg_loss),
+            'v_nb': '{}'.format(self.experiment.version),
+            'epoch': '{}'.format(self.current_epoch),
+            'batch_nb':'{}'.format(self.batch_nb),
+        }
+        tqdm_dic.update(self.tqdm_metrics)
+
+        if self.on_gpu:
+            tqdm_dic['gpu'] = '{}'.format(self.current_gpu_name)
+
+        return tqdm_dic
+
+    @property
+    def tng_tqdm_dic(self):
+        """
+        Read-only for tqdm metrics
+        :return:
+        """
+        return self.__tng_tqdm_dic
+
+    def __layout_bookeeping(self):
+
+        # determine number of training batches
+        self.nb_tng_batches = len(self.tng_dataloader)
+        self.nb_tng_batches = int(self.nb_tng_batches * self.train_percent_check)
+
+        # determine number of validation batches
+        self.nb_val_batches = len(self.val_dataloader)
+        self.nb_val_batches = int(self.nb_val_batches * self.val_percent_check)
+        self.nb_val_batches = max(1, self.nb_val_batches)
+        self.nb_val_batches = self.nb_val_batches
+
+        # determine number of test batches
+        self.nb_test_batches = len(self.test_dataloader)
+        self.nb_test_batches = int(self.nb_test_batches * self.test_percent_check)
+
+        # determine when to check validation
+        self.val_check_batch = int(self.nb_tng_batches * self.val_check_interval)
+
+    def __add_tqdm_metrics(self, metrics):
+        for k, v in metrics.items():
+            if type(v) is torch.Tensor:
+                v = v.item()
+
+            self.tqdm_metrics[k] = v
+
+    def validate(self, model, dataloader, max_batches):
+        """
+        Run validation code
+        :param model: PT model
+        :param dataloader: PT dataloader
+        :param max_batches: Scalar
+        :return:
+        """
+        # enable eval mode
+        model.zero_grad()
+        model.eval()
+
+        # disable gradients to save memory
+        torch.set_grad_enabled(False)
+
+        # bookkeeping
+        outputs = []
+
+        # run training
+        for batch_i, data_batch in enumerate(dataloader):
+
+            if data_batch is None:  # pragma: no cover
+                continue
+
+            # stop short when on fast dev run
+            if max_batches is not None and batch_i >= max_batches:
+                break
+
+            # -----------------
+            # RUN VALIDATION STEP
+            # -----------------
+            if self.use_ddp:
+                output = model(data_batch, batch_i)
+            elif self.use_dp:
+                output = model(data_batch, batch_i)
+                output = reduce_distributed_output(output, len(self.data_parallel_device_ids))
+
+            else:
+                output = model.validation_step(data_batch, batch_i)
+
+            outputs.append(output)
+
+            # batch done
+            if self.progress_bar and self.prog_bar is not None:
+                self.prog_bar.update(1)
+
+        # give model a chance to do something with the outputs
+        if self.data_parallel:
+            val_results = model.module.validation_end(outputs)
+        else:
+            val_results = model.validation_end(outputs)
+
+        # enable train mode again
+        model.train()
+
+        # enable gradients to save memory
+        torch.set_grad_enabled(True)
+
+        return val_results
+
+    def get_dataloaders(self, model):
+        """
+        Dataloaders are provided by the model
+        :param model:
+        :return:
+        """
+        self.tng_dataloader = model.tng_dataloader
+        self.test_dataloader = model.test_dataloader
+        self.val_dataloader = model.val_dataloader
+
+        if self.use_ddp and not isinstance(self.tng_dataloader.sampler, DistributedSampler):
+            msg = '''
+            when using multiple gpus and multiple nodes you must pass a DistributedSampler to DataLoader(sampler).
+            
+            ie: this:
+            dataset = myDataset()
+            dataloader = Dataloader(dataset)
+            
+            becomes:
+            dataset = myDataset()
+            dist_sampler = torch.utils.data.distributed.DistributedSampler(dataset)
+            dataloader = Dataloader(dataset, sampler=dist_sampler)
+            '''
+            raise MisconfigurationException(msg)
+
+    # -----------------------------
+    # MODEL TRAINING
+    # -----------------------------
+    def fit(self, model):
+
+        # when using multi-node or DDP within a node start each module in a separate process
+        if self.use_ddp:
+            # must copy only the meta of the exp so it survives pickle/unpickle when going to new process
+            self.experiment = self.experiment.get_meta_copy()
+
+            if self.is_slurm_managing_tasks:
+                task = int(os.environ['SLURM_LOCALID'])
+                self.ddp_train(task, model)
+            else:
+                msg = f"""
+                You requested {self.nb_requested_gpus} GPUs but launched {self.nb_slurm_tasks} slurm tasks. 
+                We will launch {self.nb_requested_gpus} processes for you. 
+                We recommend you let slurm manage the processes by setting: --ntasks-per-node={self.nb_requested_gpus}
+                If you're not using SLURM, ignore this message!
+                """
+                warnings.warn(msg)
+                mp.spawn(self.ddp_train, nprocs=len(self.data_parallel_device_ids), args=(model, ))
+
+        # 1 gpu or dp option triggers training using DP module
+        # easier to avoid NCCL issues
+        elif self.use_dp:
+            self.__dp_train(model)
+
+        # ON CPU
+        else:
+            # run through amp wrapper
+            if self.use_amp:
+                raise MisconfigurationException('amp + cpu is not supported. Please use a GPU option')
+
+            # CHOOSE OPTIMIZER
+            # filter out the weights that were done on gpu so we can load on good old cpus
+            self.optimizers = model.configure_optimizers()
+
+            self.__run_pretrain_routine(model)
+
+        # return 1 when finished
+        # used for testing or when we need to know that training succeeded
+        return 1
+
+    def __dp_train(self, model):
+
+        # CHOOSE OPTIMIZER
+        # filter out the weights that were done on gpu so we can load on good old cpus
+        self.optimizers = model.configure_optimizers()
+
+        model.cuda(self.data_parallel_device_ids[0])
+
+        # check for this bug (amp + dp + !01 doesn't work)
+        # https://github.com/NVIDIA/apex/issues/227
+        if self.use_dp and self.use_amp:
+            m = f'amp level {self.amp_level} with DataParallel is not supported. ' \
+                f'See this note from NVIDIA for more info: https://github.com/NVIDIA/apex/issues/227. ' \
+                f'We recommend you switch to ddp if you want to use amp'
+            raise MisconfigurationException(m)
+
+        model = LightningDataParallel(model, device_ids=self.data_parallel_device_ids)
+
+        self.__run_pretrain_routine(model)
+
+    def ddp_train(self, gpu_nb, model):
+        """
+        Entry point into a DP thread
+        :param gpu_nb:
+        :param model:
+        :param cluster_obj:
+        :return:
+        """
+        # node rank using relative slurm id
+        # otherwise default to node rank 0
+        try:
+            node_id = os.environ['SLURM_NODEID']
+            self.node_rank = int(node_id)
+        except Exception as e:
+            self.node_rank = 0
+
+        # recover original exp before went into process
+        # init in write mode only on proc 0
+        self.experiment.debug = self.proc_rank > 0
+        self.experiment = self.experiment.get_non_ddp_exp()
+
+        # show progbar only on prog_rank 0
+        self.prog_bar = self.prog_bar and self.node_rank == 0 and gpu_nb == 0
+
+        # determine which process we are and world size
+        self.proc_rank = self.node_rank * len(self.data_parallel_device_ids) + gpu_nb
+        self.world_size = self.nb_gpu_nodes * len(self.data_parallel_device_ids)
+
+        # let the exp know the rank to avoid overwriting logs
+        self.experiment.rank = self.proc_rank
+
+        # set up server using proc 0's ip address
+        # try to init for 20 times at max in case ports are taken
+        # where to store ip_table
+        self.__init_tcp_connection()
+
+        # CHOOSE OPTIMIZER
+        # filter out the weights that were done on gpu so we can load on good old cpus
+        self.optimizers = model.configure_optimizers()
+
+        # MODEL
+        # copy model to each gpu
+        torch.cuda.set_device(gpu_nb)
+        model.cuda(gpu_nb)
+
+        # AMP
+        # run through amp wrapper before going to distributed DP
+        if self.use_amp:
+            # An example
+            model, optimizers = amp.initialize(
+                model, self.optimizers, opt_level=self.amp_level,
+            )
+            self.optimizers = optimizers
+
+        model = LightningDistributedDataParallel(model, device_ids=[gpu_nb], find_unused_parameters=True)
+
+        # continue training routine
+        self.__run_pretrain_routine(model)
+
+    def __init_tcp_connection(self):
+        """
+        Connect all procs in the world using the env:// init
+        Use the first node as the root address
+        :param port:
+        :param tries:
+        :return:
+        """
+        # sets the appropriate port
+        try:
+            port = os.environ['MASTER_PORT']
+        except Exception as e:
+            port = 12910
+            os.environ['MASTER_PORT'] = f'{port}'
+
+        # figure out the root node addr
+        try:
+            root_node = os.environ['SLURM_NODELIST'].split(' ')[0]
+        except Exception as e:
+            root_node = '127.0.0.2'
+
+        root_node = self.resolve_root_node_address(root_node)
+        os.environ['MASTER_ADDR'] = root_node
+
+        dist.init_process_group("nccl", rank=self.proc_rank, world_size=self.world_size)
+
+    def resolve_root_node_address(self, root_node):
+        if '[' in root_node:
+            name = root_node.split('[')[0]
+            number = root_node.split(',')[0]
+            if '-' in number:
+                number = number.split('-')[0]
+
+            number = re.sub('[^0-9]', '', number)
+            root_node = name + number
+
+        return root_node
+
+    def __run_pretrain_routine(self, model):
+        """
+        Sanity check a few things before starting actual training
+        :param model:
+        :return:
+        """
+        ref_model = model
+        if self.data_parallel:
+            ref_model = model.module
+
+        ref_model.trainer = self
+
+        # set local properties on the model
+        ref_model.on_gpu = self.on_gpu
+
+        # transfer data loaders from model
+        self.get_dataloaders(ref_model)
+
+        # init training constants
+        self.__layout_bookeeping()
+
+        # add lr schedulers
+        if self.lr_scheduler_milestones is not None:
+            for optimizer in self.optimizers:
+                scheduler = MultiStepLR(optimizer, self.lr_scheduler_milestones)
+                self.lr_schedulers.append(scheduler)
+
+        # print model summary
+        if self.proc_rank == 0 and self.print_weights_summary:
+            ref_model.summarize()
+
+        # give model convenience properties
+        ref_model.trainer = self
+        ref_model.experiment = self.experiment
+
+        # run tiny validation to make sure program won't crash during val
+        _ = self.validate(model, self.val_dataloader, max_batches=self.nb_sanity_val_steps)
+
+        # save exp to get started
+        if self.proc_rank == 0:
+            self.experiment.save()
+
+        # track model now.
+        # if cluster resets state, the model will update with the saved weights
+        self.model = model
+
+        # enable cluster checkpointing
+        # also restores training state
+        if self.cluster is not None:  # pragma: no cover
+            self.enable_auto_hpc_walltime_manager()
+
+        # ---------------------------
+        # CORE TRAINING LOOP
+        # ---------------------------
+        self.__train()
+
+    def __train(self):
+        # run all epochs
+        for epoch_nb in range(self.current_epoch, self.max_nb_epochs):
+            # update the lr scheduler
+            for lr_scheduler in self.lr_schedulers:
+                lr_scheduler.step()
+
+            model = self.__get_model()
+            model.current_epoch = epoch_nb
+
+            # hook
+            if self.__is_function_implemented('on_epoch_start'):
+                model = self.__get_model()
+                model.on_epoch_start()
+
+            self.current_epoch = epoch_nb
+            self.total_batches = self.nb_tng_batches + self.nb_val_batches
+            self.batch_loss_value = 0  # accumulated grads
+
+            # init progbar when requested
+            if self.progress_bar:
+                self.prog_bar = tqdm.tqdm(range(self.total_batches), position=self.process_position)
+
+            for batch_nb, data_batch in enumerate(self.tng_dataloader):
+                self.batch_nb = batch_nb
+                self.global_step += 1
+
+                model = self.__get_model()
+                model.global_step = self.global_step
+
+                # stop when the flag is changed or we've gone past the amount requested in the batches
+                self.total_batch_nb += 1
+                met_batch_limit = batch_nb > self.nb_tng_batches
+                if met_batch_limit:
+                    break
+
+                # ---------------
+                # RUN TRAIN STEP
+                # ---------------
+                batch_result = self.__run_tng_batch(data_batch, batch_nb)
+                early_stop_epoch = batch_result == -1
+
+                # ---------------
+                # RUN VAL STEP
+                # ---------------
+                is_val_check_batch = (batch_nb + 1) % self.val_check_batch == 0
+                if self.fast_dev_run or is_val_check_batch or early_stop_epoch:
+                    self.__run_validation()
+
+                # when batch should be saved
+                if (batch_nb + 1) % self.log_save_interval == 0 or early_stop_epoch:
+                    if self.proc_rank == 0:
+                        self.experiment.save()
+
+                # when metrics should be logged
+                if batch_nb % self.add_log_row_interval == 0 or early_stop_epoch:
+                    # count items in memory
+                    # nb_params, nb_tensors = count_mem_items()
+
+                    model = self.__get_model()
+                    metrics = self.__tng_tqdm_dic
+
+                    # add gpu memory
+                    if self.on_gpu:
+                        mem_map = get_gpu_memory_map()
+                        metrics.update(mem_map)
+
+                    # add norms
+                    if self.track_grad_norm > 0:
+                        model = self.__get_model()
+                        grad_norm_dic = model.grad_norm(self.track_grad_norm)
+                        metrics.update(grad_norm_dic)
+
+                    if self.__is_function_implemented('on_tng_metrics'):
+                        model.on_tng_metrics(metrics)
+
+                    # log metrics
+                    scalar_metrics = self.__metrics_to_scalars(metrics, blacklist=self.__log_vals_blacklist())
+                    if self.proc_rank == 0:
+                        self.experiment.log(scalar_metrics, global_step=self.global_step)
+                        self.experiment.save()
+
+                # hook
+                if self.__is_function_implemented('on_batch_end'):
+                    model = self.__get_model()
+                    model.on_batch_end()
+
+                # end epoch early
+                if early_stop_epoch:
+                    break
+
+            # hook
+            if self.__is_function_implemented('on_epoch_end'):
+                model = self.__get_model()
+                model.on_epoch_end()
+
+            # early stopping
+            met_min_epochs = epoch_nb > self.min_nb_epochs
+            if self.enable_early_stop and met_min_epochs:
+                should_stop = self.early_stop_callback.on_epoch_end(epoch=epoch_nb, logs=self.__tng_tqdm_dic)
+
+                # stop training
+                stop = should_stop and met_min_epochs
+                if stop:
+                    return
+
+    def __metrics_to_scalars(self, metrics, blacklist=[]):
+        new_metrics = {}
+        for k, v in metrics.items():
+            if type(v) is torch.Tensor:
+                v = v.item()
+
+            if type(v) is dict:
+                v = self.__metrics_to_scalars(v)
+
+            if k not in blacklist:
+                new_metrics[k] = float(v)
+
+        return new_metrics
+
+    def __log_vals_blacklist(self):
+        """avoid logging some vals lightning uses to maintain state"""
+        blacklist = {'batch_nb', 'v_nb', 'gpu'}
+        return blacklist
+
+    def __run_tng_batch(self, data_batch, batch_nb):
+        if data_batch is None:
+            return 0
+
+        # hook
+        if self.__is_function_implemented('on_batch_start'):
+            model_ref = self.__get_model()
+            response = model_ref.on_batch_start(data_batch)
+
+            if response == -1:
+                return -1
+
+        if self.progress_bar:
+            self.prog_bar.update(1)
+
+        # forward pass
+        # return a scalar value and a dic with tqdm metrics
+        if self.use_ddp:
+            output = self.model(data_batch, batch_nb)
+        elif self.use_dp:
+            output = self.model(data_batch, batch_nb)
+            output = reduce_distributed_output(output, len(self.data_parallel_device_ids))
+        else:
+            output = self.model.training_step(data_batch, batch_nb)
+
+        try:
+            model_specific_tqdm_metrics_dic = output['tqdm_metrics']
+        except Exception as e:
+            model_specific_tqdm_metrics_dic = {}
+
+        # if output dict doesn't have the keyword loss
+        # then assume the output=loss if scalar
+        try:
+            loss = output['loss']
+        except Exception as e:
+            if type(output) is torch.Tensor:
+                loss = output
+
+        self.__add_tqdm_metrics(model_specific_tqdm_metrics_dic)
+
+        # backward pass
+        if self.use_amp:
+            # scale loss when using amp
+            for optimizer in self.optimizers:
+                with amp.scale_loss(loss, optimizer) as scaled_loss:
+                    scaled_loss.backward()
+        else:
+            loss.backward()
+
+        # insert after step hook
+        if self.__is_function_implemented('on_after_backward'):
+            model_ref = self.__get_model()
+            response = model_ref.on_after_backward()
+
+        if self.print_nan_grads:
+            model = self.__get_model()
+            for param in model.parameters():
+                print(param.grad.float().sum())
+
+        # avoid memory leaks
+        self.batch_loss_value += loss.item()
+
+        # gradient update with accumulated gradients
+        if (self.batch_nb + 1) % self.accumulate_grad_batches == 0:
+
+            # clip gradients
+            if self.gradient_clip > 0:
+                model = self.__get_model()
+                torch.nn.utils.clip_grad_norm(model.parameters(), self.gradient_clip)
+
+            # update gradients across all optimizers
+            for optimizer in self.optimizers:
+                optimizer.step()
+
+                # insert after step hook
+                if self.__is_function_implemented('on_before_zero_grad'):
+                    model_ref = self.__get_model()
+                    response = model_ref.on_before_zero_grad(optimizer)
+
+                # clear gradients
+                optimizer.zero_grad()
+
+            # queuing loss across batches blows it up proportionally... divide out the number accumulated
+            self.batch_loss_value = self.batch_loss_value / self.accumulate_grad_batches
+
+            # track loss
+            self.running_loss.append(self.batch_loss_value)
+            self.batch_loss_value = 0
+            self.avg_loss = np.mean(self.running_loss[-100:])
+
+            # update progbar
+            if self.progress_bar:
+                # add model specific metrics
+                tqdm_metrics = self.__tng_tqdm_dic
+                self.prog_bar.set_postfix(**tqdm_metrics)
+
+        # activate batch end hook
+        if self.__is_function_implemented('on_batch_end'):
+            model = self.__get_model()
+            model.on_batch_end()
+
+        return 0
+
+    def __run_validation(self):
+        # decide if can check epochs
+        can_check_epoch = (self.current_epoch + 1) % self.check_val_every_n_epoch == 0
+        if self.fast_dev_run:
+            print('skipping to check performance bc of --fast_dev_run')
+        elif not can_check_epoch:
+            return
+
+        # hook
+        if self.__is_function_implemented('on_pre_performance_check'):
+            model = self.__get_model()
+            model.on_pre_performance_check()
+
+        # use full val set on end of epoch
+        # use a small portion otherwise
+        max_batches = None if not self.fast_dev_run else 1
+        model_specific_tqdm_metrics_dic = self.validate(
+            self.model,
+            self.val_dataloader,
+            max_batches
+        )
+        self.__add_tqdm_metrics(model_specific_tqdm_metrics_dic)
+
+        # hook
+        if self.__is_function_implemented('on_post_performance_check'):
+            model = self.__get_model()
+            model.on_post_performance_check()
+
+        if self.progress_bar:
+            # add model specific metrics
+            tqdm_metrics = self.__tng_tqdm_dic
+            self.prog_bar.set_postfix(**tqdm_metrics)
+
+        # model checkpointing
+        if self.proc_rank == 0 and self.checkpoint_callback is not None:
+            print('save callback...')
+            self.checkpoint_callback.on_epoch_end(epoch=self.current_epoch, logs=self.__tng_tqdm_dic)
@@ -1,6 +1,7 @@
 from torch.nn import DataParallel
 from torch.nn.parallel import DistributedDataParallel
 import itertools
+from itertools import chain

 import threading
 import torch
@@ -8,7 +9,7 @@ from torch.cuda._utils import _get_device_index
 import pdb


-def _find_tensors(obj):
+def _find_tensors(obj):  # pragma: no cover
    r"""
    Recursively find all tensors contained in the specified object.
    """
@@ -21,8 +22,7 @@ def _find_tensors(obj):
    return []


-
-def get_a_var(obj):
+def get_a_var(obj):  # pragma: no cover
    if isinstance(obj, torch.Tensor):
        return obj

@@ -42,6 +42,29 @@ class LightningDataParallel(DataParallel):
    Override the forward call in lightning so it goes to training and validation step respectively
    """

+    def forward(self, *inputs, **kwargs):
+        if not self.device_ids:
+            return self.module(*inputs, **kwargs)
+
+        for t in chain(self.module.parameters(), self.module.buffers()):
+            if t.device != self.src_device_obj:
+                raise RuntimeError("module must have its parameters and buffers "
+                                   "on device {} (device_ids[0]) but found one of "
+                                   "them on device: {}".format(self.src_device_obj, t.device))
+
+        inputs, kwargs = self.scatter(inputs, kwargs, self.device_ids)
+        if len(self.device_ids) == 1:
+            # lightning
+            if self.module.training:
+                return self.module.training_step(*inputs[0], **kwargs[0])
+            else:
+                return self.module.validation_step(*inputs[0], **kwargs[0])
+
+        replicas = self.replicate(self.module, self.device_ids[:len(inputs)])
+        outputs = self.parallel_apply(replicas, inputs, kwargs)
+        return self.gather(outputs, self.output_device)
+
+
    def parallel_apply(self, replicas, inputs, kwargs):
        return parallel_apply(replicas, inputs, kwargs, self.device_ids[:len(replicas)])

@@ -54,7 +77,7 @@ class LightningDistributedDataParallel(DistributedDataParallel):
    def parallel_apply(self, replicas, inputs, kwargs):
        return parallel_apply(replicas, inputs, kwargs, self.device_ids[:len(replicas)])

-    def forward(self, *inputs, **kwargs):
+    def forward(self, *inputs, **kwargs):  # pragma: no cover
        self._sync_params()
        if self.device_ids:
            inputs, kwargs = self.scatter(inputs, kwargs, self.device_ids)
@@ -89,7 +112,7 @@ class LightningDistributedDataParallel(DistributedDataParallel):
        return output


-def parallel_apply(modules, inputs, kwargs_tup=None, devices=None):
+def parallel_apply(modules, inputs, kwargs_tup=None, devices=None):  # pragma: no cover
    r"""Applies each `module` in :attr:`modules` in parallel on arguments
    contained in :attr:`inputs` (positional) and :attr:`kwargs_tup` (keyword)
    on each of :attr:`devices`.
@@ -0,0 +1,17 @@
+
+def data_loader(fn):
+    """
+    Decorator to make any fx with this use the lazy property
+    :param fn:
+    :return:
+    """
+
+    attr_name = '_lazy_' + fn.__name__
+
+    @property
+    def _data_loader(self):
+        if not hasattr(self, attr_name):
+            setattr(self, attr_name, fn(self))
+        return getattr(self, attr_name)
+
+    return _data_loader
@@ -27,14 +27,3 @@ class GradInformation(nn.Module):
        results['grad_{}_norm_total'.format(norm_type)] = round(total_norm.data.cpu().numpy().flatten()[0], 3)
        return results

-
-    def describe_grads(self):
-        for p in self.parameters():
-            g = p.grad.data.numpy().flatten()
-            print(np.max(g), np.min(g), np.mean(g))
-
-
-    def describe_params(self):
-        for p in self.parameters():
-            g = p.data.numpy().flatten()
-            print(np.max(g), np.min(g), np.mean(g))
@@ -19,3 +19,27 @@ class ModelHooks(torch.nn.Module):
    def on_post_performance_check(self):
        pass

+    def on_tng_metrics(self, metrics):
+        pass
+
+    def on_before_zero_grad(self, optimizer):
+        """
+        Called after optimizer.step() and before optimizer.zero_grad()
+
+        for optimizer in optimizers:
+            optimizer.step()
+            model.on_before_zero_grad(optimizer) # < ---- called here
+            optimizer.zero_grad
+
+        :param optimizer:
+        :return:
+        """
+        pass
+
+    def on_after_backward(self):
+        """
+        Called after loss.backward() and before optimizers do anything
+        :return:
+        """
+        pass
+
@@ -33,33 +33,42 @@ class ModelSummary(object):
        mods = list(self.model.modules())
        in_sizes = []
        out_sizes = []
-        input_ = self.example_input_array
-        for i in range(1, len(mods)):
-            m = mods[i]
-            if type(input_) is list or type(input_) is tuple:
-                out = m(*input_)
-            else:
-                out = m(input_)
+        input_ = self.model.example_input_array

-            if type(input_) is tuple or type(input_) is list:
-                in_size = []
-                for x in input_:
-                    if type(x) is list:
-                        in_size.append(len(x))
-                    else:
-                        in_size.append(x.size())
-            else:
-                in_size = np.array(input_.size())
+        if self.model.on_gpu:
+            input_ = input_.cuda(0)

-            in_sizes.append(in_size)
+        if self.model.trainer.use_amp:
+            input_ = input_.half()

-            if type(out) is tuple or type(out) is list:
-                out_size = np.asarray([x.size() for x in out])
-            else:
-                out_size = np.array(out.size())
+        with torch.no_grad():

-            out_sizes.append(out_size)
-            input_ = out
+            for i in range(1, len(mods)):
+                m = mods[i]
+                if type(input_) is list or type(input_) is tuple:  # pragma: no cover
+                    out = m(*input_)
+                else:
+                    out = m(input_)
+
+                if type(input_) is tuple or type(input_) is list:  # pragma: no cover
+                    in_size = []
+                    for x in input_:
+                        if type(x) is list:
+                            in_size.append(len(x))
+                        else:
+                            in_size.append(x.size())
+                else:
+                    in_size = np.array(input_.size())
+
+                in_sizes.append(in_size)
+
+                if type(out) is tuple or type(out) is list:  # pragma: no cover
+                    out_size = np.asarray([x.size() for x in out])
+                else:
+                    out_size = np.array(out.size())
+
+                out_sizes.append(out_size)
+                input_ = out

        self.in_sizes = in_sizes
        self.out_sizes = out_sizes
@@ -114,13 +123,22 @@ class ModelSummary(object):
        Layer Name, Layer Type, Input Size, Output Size, Number of Parameters
        '''

-        df = pd.DataFrame( np.zeros( (len(self.layer_names), 3) ) )
-        df.columns = ['Name', 'Type', 'Params']
+        cols = ['Name', 'Type', 'Params']
+        if self.model.example_input_array is not None:
+            cols.extend(['In_sizes', 'Out_sizes'])
+
+        df = pd.DataFrame(np.zeros( (len(self.layer_names), len(cols))))
+        df.columns = cols

        df['Name'] = self.layer_names
        df['Type'] = self.layer_types
        df['Params'] = self.param_nums

+        if self.model.example_input_array is not None:
+
+            df['In_sizes'] = self.in_sizes
+            df['Out_sizes'] = self.out_sizes
+
        self.summary = df
        return

@@ -128,10 +146,13 @@ class ModelSummary(object):
        self.get_layer_names()
        self.get_parameter_sizes()
        self.get_parameter_nums()
+
+        if self.model.example_input_array is not None:
+            self.get_variable_sizes()
        self.make_summary()


-def print_mem_stack():
+def print_mem_stack(): # pragma: no cover
    for obj in gc.get_objects():
        try:
            if torch.is_tensor(obj) or (hasattr(obj, 'data') and torch.is_tensor(obj.data)):
@@ -140,7 +161,7 @@ def print_mem_stack():
            pass


-def count_mem_items():
+def count_mem_items(): # pragma: no cover
    nb_params = 0
    nb_tensors = 0
    for obj in gc.get_objects():
@@ -2,36 +2,38 @@ import torch
 import os
 import re
 import pdb
-from pytorch_lightning.pt_overrides.override_data_parallel import LightningDistributedDataParallel
+from pytorch_lightning.pt_overrides.override_data_parallel import LightningDistributedDataParallel, LightningDataParallel
+

 class ModelIO(object):

-    def load_model_specific(self, checkpoint):
+    def on_load_checkpoint(self, checkpoint):
        """
        Do something with the checkpoint
+        Gives model a chance to load something before state_dict is restored
        :param checkpoint:
        :return:
        """
-        raise NotImplementedError
+        pass

-    def get_save_dict(self):
+    def on_save_checkpoint(self, checkpoint):
        """
-        Return specific things for the model
-        :return:
+        Give the model a chance to add something to the checkpoint.
+        state_dict is already there
        """
-        raise NotImplementedError
+        pass

    # -------------------------
    # OPTIONAL HOOKS
    # -------------------------
-    def on_hpc_save(self):
+    def on_hpc_save(self, checkpoint):
        """
        Hook to do whatever you need right before Slurm manager saves the model
        :return:
        """
        pass

-    def on_hpc_load(self):
+    def on_hpc_load(self, checkpoint):
        """
        Hook to do whatever you need right before Slurm manager loads the model
        :return:
@@ -41,6 +43,11 @@ class ModelIO(object):

 class TrainerIO(object):

+    def __get_model(self):
+        is_dp_module = type(self.model) is LightningDistributedDataParallel or type(self.model) is LightningDataParallel
+        model = self.model.module if is_dp_module else self.model
+        return model
+
    # --------------------
    # MODEL SAVE CHECKPOINT
    # --------------------
@@ -51,26 +58,32 @@ class TrainerIO(object):
        torch.save(checkpoint, filepath)

    def dump_checkpoint(self):
+
        checkpoint = {
            'epoch': self.current_epoch,
-            'checkpoint_callback_best': self.checkpoint_callback.best,
-            'early_stop_callback_wait': self.early_stop_callback.wait,
-            'early_stop_callback_patience': self.early_stop_callback.patience,
            '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
+
        optimizer_states = []
        for i, optimizer in enumerate(self.optimizers):
            optimizer_states.append(optimizer.state_dict())

        checkpoint['optimizer_states'] = optimizer_states

-        # request what to save from the model
-        model = self.model.module if type(self.model) is LightningDistributedDataParallel else self.model
-        checkpoint_dict = model.get_save_dict()
+        # 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)

-        # merge trainer and model saving items
-        checkpoint.update(checkpoint_dict)
        return checkpoint

    # --------------------
@@ -103,9 +116,13 @@ class TrainerIO(object):
        :param checkpoint:
        :return:
        """
-        self.checkpoint_callback.best = checkpoint['checkpoint_callback_best']
-        self.early_stop_callback.wait = checkpoint['early_stop_callback_wait']
-        self.early_stop_callback.patience = checkpoint['early_stop_callback_patience']
+        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']

@@ -134,13 +151,15 @@ class TrainerIO(object):
        filepath = '{}/hpc_ckpt_{}.ckpt'.format(folderpath, ckpt_number)

        # give model a chance to do something on hpc_save
-        self.on_hpc_save()
+        model = self.__get_model()
+        checkpoint = self.dump_checkpoint()

-        # request what to save from the model
-        checkpoint_dict = self.dump_checkpoint()
+        model.on_hpc_save(checkpoint)

        # do the actual save
-        torch.save(checkpoint_dict, filepath)
+        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))
@@ -150,15 +169,17 @@ class TrainerIO(object):
        else:
            checkpoint = torch.load(filepath, map_location=lambda storage, loc: storage)

-        # load training state
+        # load training state (affects trainer only)
        self.restore_training_state(checkpoint)

        # load model state
-        model = self.model.module if type(self.model) is LightningDataParallel else self.model
-        model.load_model_specific(checkpoint)
+        model = self.__get_model()
+
+        # load the state_dict on the model automatically
+        model.load_state_dict(checkpoint['state_dict'])

        # call model hook
-        self.on_hpc_load()
+        model.on_hpc_load(checkpoint)

    def max_ckpt_in_folder(self, path):
        files = os.listdir(path)
@@ -1,22 +0,0 @@
-from torch import nn
-from torch import optim
-
-
-class OptimizerConfig(nn.Module):
-
-    def choose_optimizer(self, optimizer, params, optimizer_params, opt_name_key):
-        if optimizer == 'adam':
-            optimizer = optim.Adam(params, **optimizer_params)
-        if optimizer == 'sparse_adam':
-            optimizer = optim.SparseAdam(params, **optimizer_params)
-        if optimizer == 'sgd':
-            optimizer = optim.SGD(params, **optimizer_params)
-        if optimizer == 'adadelta':
-            optimizer = optim.Adadelta(params, **optimizer_params)
-
-        # transfer opt state if loaded
-        if opt_name_key in self.loaded_optimizer_states_dict:
-            state = self.loaded_optimizer_states_dict[opt_name_key]
-            optimizer.load_state_dict(state)
-
-        return optimizer
@@ -1,19 +1,15 @@
-import os
 import torch
-import math
-
 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.root_module.optimization import OptimizerConfig
 from pytorch_lightning.root_module.hooks import ModelHooks
+from pytorch_lightning.root_module.decorators import data_loader


-class LightningModule(GradInformation, ModelIO, OptimizerConfig, ModelHooks):
+class LightningModule(GradInformation, ModelIO, ModelHooks):

-    def __init__(self, hparams):
-        super(LightningModule, self).__init__()
-        self.hparams = hparams
+    def __init__(self, *args, **kwargs):
+        super(LightningModule, self).__init__(*args, **kwargs)

        self.dtype = torch.FloatTensor
        self.exp_save_path = None
@@ -22,15 +18,11 @@ class LightningModule(GradInformation, ModelIO, OptimizerConfig, ModelHooks):
        self.loaded_optimizer_states_dict = {}
        self.trainer = None
        self.experiment = None
+        self.example_input_array = None

        # track if gpu was requested for checkpointing
        self.on_gpu = False

-        # computed vars for the dataloaders
-        self._tng_dataloader = None
-        self._val_dataloader = None
-        self._test_dataloader = None
-
    def forward(self, *args, **kwargs):
        """
        Expand model in into whatever you need.
@@ -71,37 +63,7 @@ class LightningModule(GradInformation, ModelIO, OptimizerConfig, ModelHooks):
        """
        raise NotImplementedError

-    def update_tng_log_metrics(self, logs):
-        """
-        Chance to update metrics to be logged for training step.
-        For example, add music, images, etc... to log
-        :param logs:
-        :return:
-        """
-        raise NotImplementedError
-
-    def loss(self, *args, **kwargs):
-        """
-        Expand model_out into your components
-        :param model_out:
-        :return:
-        """
-        raise NotImplementedError
-
-    def summarize(self):
-        model_summary = ModelSummary(self)
-        print(model_summary)
-
-
-    def freeze(self):
-        for param in self.parameters():
-            param.requires_grad = False
-
-    def unfreeze(self):
-        for param in self.parameters():
-            param.requires_grad = True
-
-    @property
+    @data_loader
    def tng_dataloader(self):
        """
        Implement a function to load an h5py of this data
@@ -109,7 +71,7 @@ class LightningModule(GradInformation, ModelIO, OptimizerConfig, ModelHooks):
        """
        raise NotImplementedError

-    @property
+    @data_loader
    def test_dataloader(self):
        """
        Implement a function to load an h5py of this data
@@ -117,7 +79,7 @@ class LightningModule(GradInformation, ModelIO, OptimizerConfig, ModelHooks):
        """
        raise NotImplementedError

-    @property
+    @data_loader
    def val_dataloader(self):
        """
        Implement a function to load an h5py of this data
@@ -125,16 +87,6 @@ class LightningModule(GradInformation, ModelIO, OptimizerConfig, ModelHooks):
        """
        raise NotImplementedError

-    @staticmethod
-    def get_process_position(gpus):
-        try:
-            current_gpu = os.environ["CUDA_VISIBLE_DEVICES"]
-            gpu_ids = gpus.split(';')
-            process_position = gpu_ids.index(current_gpu)
-            return process_position, current_gpu
-        except Exception as e:
-            return 0, 0
-
    @classmethod
    def load_from_metrics(cls, weights_path, tags_csv, on_gpu, map_location=None):
        """
@@ -156,9 +108,26 @@ class LightningModule(GradInformation, ModelIO, OptimizerConfig, ModelHooks):
        else:
            checkpoint = torch.load(weights_path, map_location=lambda storage, loc: storage)

+        # load the state_dict on the model automatically
        model = cls(hparams)
+        model.load_state_dict(checkpoint['state_dict'])
+
+        # give model a chance to load something
+        model.on_load_checkpoint(checkpoint)

-        # allow model to load
-        model.load_model_specific(checkpoint)
-        model.load_state_dict(checkpoint['state_dict'], strict=False)
        return model
+
+    def summarize(self):
+        model_summary = ModelSummary(self)
+        print(model_summary)
+
+    def freeze(self):
+        for param in self.parameters():
+            param.requires_grad = False
+
+    def unfreeze(self):
+        for param in self.parameters():
+            param.requires_grad = True
+
+
+
@@ -0,0 +1,253 @@
+import os
+from collections import OrderedDict
+import torch.nn as nn
+from torchvision.datasets import MNIST
+import torchvision.transforms as transforms
+import torch
+import torch.nn.functional as F
+from test_tube import HyperOptArgumentParser
+from torch import optim
+from torch.utils.data import DataLoader
+from torch.utils.data.distributed import DistributedSampler
+
+from pytorch_lightning.root_module.root_module import LightningModule
+import pytorch_lightning as ptl
+
+
+class LightningTestModel(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(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)
+
+        output = OrderedDict({
+            'loss': loss_val
+        })
+
+        # can also return just a scalar instead of a dict (return loss_val)
+        return output
+
+    def validation_step(self, data_batch, batch_i):
+        """
+        Lightning calls this inside the validation loop
+        :param data_batch:
+        :return:
+        """
+        x, y = data_batch
+        x = x.view(x.size(0), -1)
+        y_hat = self.forward(x)
+
+        loss_val = self.loss(y, y_hat)
+
+        # acc
+        labels_hat = torch.argmax(y_hat, dim=1)
+        val_acc = torch.sum(y == labels_hat).item() / (len(y) * 1.0)
+        val_acc = torch.tensor(val_acc)
+
+        if self.on_gpu:
+            val_acc = val_acc.cuda(loss_val.device.index)
+
+        # in DP mode (default) make sure if result is scalar, there's another dim in the beginning
+        if self.trainer.use_dp:
+            loss_val = loss_val.unsqueeze(0)
+            val_acc = val_acc.unsqueeze(0)
+
+        # alternate possible outputs to test
+        if self.trainer.batch_nb % 1 == 0:
+            output = OrderedDict({
+                'val_loss': loss_val,
+                'val_acc': val_acc,
+            })
+            return output
+        if self.trainer.batch_nb % 2 == 0:
+            return val_acc
+
+        if self.trainer.batch_nb % 3 == 0:
+            output = OrderedDict({
+                'val_loss': loss_val,
+                'val_acc': val_acc,
+                'test_dic': {'val_loss_a': loss_val}
+            })
+            return output
+
+    def validation_end(self, outputs):
+        """
+        Called at the end of validation to aggregate outputs
+        :param outputs: list of individual outputs of each validation step
+        :return:
+        """
+        # if returned a scalar from validation_step, outputs is a list of tensor scalars
+        # we return just the average in this case (if we want)
+        # return torch.stack(outputs).mean()
+
+        val_loss_mean = 0
+        val_acc_mean = 0
+        for output in outputs:
+            val_loss_mean += output['val_loss']
+            val_acc_mean += output['val_acc']
+
+        val_loss_mean /= len(outputs)
+        val_acc_mean /= len(outputs)
+
+        tqdm_dic = {'val_loss': val_loss_mean.item(), 'val_acc': val_acc_mean.item()}
+        return tqdm_dic
+
+    def on_tng_metrics(self, logs):
+        logs['some_tensor_to_test'] = torch.rand(1)
+
+    # ---------------------
+    # TRAINING SETUP
+    # ---------------------
+    def configure_optimizers(self):
+        """
+        return whatever optimizers we want here
+        :return: list of optimizers
+        """
+        optimizer = optim.Adam(self.parameters(), lr=self.hparams.learning_rate)
+        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.on_gpu 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 as e:
+            pass
+
+        should_shuffle = train_sampler is None
+        loader = DataLoader(
+            dataset=dataset,
+            batch_size=batch_size,
+            shuffle=should_shuffle,
+            sampler=train_sampler
+        )
+
+        return loader
+
+    @ptl.data_loader
+    def tng_dataloader(self):
+        return self.__dataloader(train=True)
+
+    @ptl.data_loader
+    def val_dataloader(self):
+        return self.__dataloader(train=False)
+
+    @ptl.data_loader
+    def test_dataloader(self):
+        return self.__dataloader(train=False)
+
+    @staticmethod
+    def add_model_specific_args(parent_parser, root_dir):
+        """
+        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)
+        parser.add_argument('--hidden_dim', default=50000, type=int) # use 500 for CPU, 50000 for GPU to see speed difference
+
+        # 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 the gpus being used across all nodes')
+        return parser
@@ -52,10 +52,6 @@ def main(hparams, cluster, results_dict):
    hparams.__setattr__('nb_gpus', torch.cuda.device_count())
    hparams.__setattr__('inference_mode', hparams.model_load_weights_path is not None)

-    # delay each training start to not overwrite logs
-    process_position, current_gpu = TRAINING_MODEL.get_process_position(hparams.gpus)
-    sleep(process_position + 1)
-
    # init experiment
    exp = Experiment(
        name=hparams.tt_name,
@@ -0,0 +1,5 @@
+import pdb
+import sys
+
+class MisconfigurationException(Exception):
+    pass
@@ -1,104 +0,0 @@
-import torch
-import numpy as np
-from copy import deepcopy
-
-
-class PretrainedEmbedding(torch.nn.Embedding):
-
-    def __init__(self, embedding_path, embedding_dim, task_vocab, freeze=True, *args, **kwargs):
-        """
-        Loads a prebuilt pytorch embedding from any embedding formated file.
-        Padding=0 by default.
-
-        >>> emb = PretrainedEmbedding(embedding_path='glove.840B.300d.txt',embedding_dim=300, task_vocab={'hello': 1, 'world': 2})
-        >>> data = torch.Tensor([[0, 1], [0, 2]]).long()
-        >>> embedded = emb(data)
-        
-
-
-        :param embedding_path:
-        :param emb_dim:
-        :param task_vocab:
-        :param freeze:
-        :return:
-        """
-        # count the vocab
-        self.vocab_size = max(task_vocab.values()) + 1
-        super(PretrainedEmbedding, self).__init__(self.vocab_size, embedding_dim, padding_idx=0, *args, **kwargs)
-
-        # load pretrained embeddings
-        new_emb = self.__load_task_specific_embeddings(deepcopy(task_vocab), embedding_path, embedding_dim, freeze)
-
-        # transfer weights
-        self.weight = new_emb.weight
-
-        # apply freeze
-        should_freeze = not freeze
-        self.weight.requires_grad = should_freeze
-
-    def __load_task_specific_embeddings(self, vocab_words, embedding_path, emb_dim, freeze):
-        """
-        Iterates embedding file to only pull out task specific embeddings
-        :param vocab_words:
-        :param embedding_path:
-        :param emb_dim:
-        :param freeze:
-        :return:
-        """
-
-        # holds final embeddings for relevant words
-        embeddings = np.zeros(shape=(self.vocab_size, emb_dim))
-
-        # load embedding line by line and extract relevant embeddings
-        with open(embedding_path, encoding='utf-8') as f:
-            for line in f:
-                tokens = line.split(' ')
-                word = tokens[0]
-                embedding = tokens[1:]
-                embedding[-1] = embedding[-1][:-1]  # remove last new line
-
-                if word in vocab_words:
-                    vocab_word_i = vocab_words[word]
-
-                    # skip words that try to overwrite pad idx
-                    if vocab_word_i == 0:
-                        del vocab_words[word]
-                        continue
-
-                    emb_vals = np.asarray([float(x) for x in embedding])
-                    embeddings[vocab_word_i] = emb_vals
-
-                    # remove vocab word to early terminate
-                    del vocab_words[word]
-
-                # early break
-                if len(vocab_words) == 0:
-                    break
-
-        # add random vectors for the non-pretrained words
-        # these are vocab words NOT found in the pretrained embeddings
-        for w, i in vocab_words.items():
-            # skip words that try to overwrite pad idx
-            if i == 0:
-                continue
-
-            embedding = np.random.normal(size=emb_dim)
-            embeddings[i] = embedding
-
-        # turn into pt embedding
-        embeddings = torch.FloatTensor(embeddings)
-        embeddings = torch.nn.Embedding.from_pretrained(embeddings, freeze=freeze)
-
-        return embeddings
-
-
-if __name__ == '__main__':
-    emb = PretrainedEmbedding(
-        embedding_path='/Users/waf/Developer',
-        embedding_dim=300,
-        task_vocab={'hello': 1, 'world': 2}
-    )
-
-    data = torch.Tensor([[0, 1], [0, 2]]).long()
-    embedded = emb(data)
-    print(embedded)
@@ -1,28 +0,0 @@
-from matplotlib import pyplot as plt
-import numpy as np
-np.seterr(divide='ignore', invalid='ignore')
-
-
-def plot_confusion_matrix(cm,
-                          save_path,
-                          normalize=False,
-                          title='Confusion matrix',
-                          ylabel='y',
-                          xlabel='x'):
-    """
-    This function prints and plots the confusion matrix.
-    Normalization can be applied by setting `normalize=True`.
-    """
-    if normalize:
-        cm = cm.astype('float') / cm.sum(axis=1)[:, np.newaxis]
-        print("Normalized confusion matrix")
-    else:
-        print('Confusion matrix, without normalization')
-
-    fig = plt.figure()
-    plt.matshow(cm)
-    plt.title(title)
-    plt.colorbar()
-    plt.ylabel(ylabel)
-    plt.xlabel(xlabel)
-    plt.savefig(save_path)
@@ -1,35 +1,9 @@
-
-atomicwrites==1.2.1
-attrs==18.2.0
-certifi==2018.11.29
-cffi==1.11.5
-imageio==2.4.1
-mkl-fft==1.0.6
-mkl-random==1.0.2
-more-itertools==5.0.0
-numpy==1.15.4
-olefile==0.46
-pandas==0.23.4
-Pillow==5.3.0
-pluggy==0.8.0
-py==1.7.0
-pycparser==2.19
-pytest==4.0.2
-python-dateutil==2.7.5
-pytz==2018.7
+coverage==4.5.3
+mkdocs==1.0.4
+pytest==5.0.1
 scikit-learn==0.20.2
-scipy==1.2.0
-six==1.12.0
-sklearn==0.0
-tensorboard==1.14.0
-tensorboardX==1.7
-tensorflow==1.14.0
-test-tube==0.643
-torch==1.0.0
-torchvision==0.2.1
 tqdm==4.32.1
 twine==1.13.0
-urllib3==1.25.3
-webencodings==0.5.1
-Werkzeug==0.15.4
-wrapt==1.11.2
+numpy==1.16.4
+torch>=1.1.0
+torchvision==0.3.0
@@ -16,6 +16,33 @@ markers =
 ignore = E731,W504
 max-line-length = 120

+[coverage:report]
+exclude_lines =
+    pragma: no cover
+    def __repr__
+    if self.debug:
+    if settings.DEBUG
+    raise AssertionError
+    raise NotImplementedError
+    if 0:
+    if __name__ == .__main__.:
+    except Exception as e
+    print(e)
+    print(traceback.print_exc())
+    return *
+    raise Exception
+    warnings
+    print
+    raise RuntimeError
+    break
+    pass
+    os.makedirs
+
+omit =
+    pytorch_lightning/callbacks/pt_callbacks.py
+    tests/test_models.py
+    pytorch_lightning/testing_models/lm_test_module.py
+
 [flake8]
 ignore = E731,W504,F401,F841
 max-line-length = 120
@@ -7,7 +7,7 @@ from setuptools import setup, find_packages
 # http://blog.ionelmc.ro/2014/05/25/python-packaging/
 setup(
    name="pytorch-lightning",
-    version='0.21',
+    version='0.3.6.4',
    description="The Keras for ML researchers using PyTorch",
    author="William Falcon",
    author_email="waf2107@columbia.edu",
@@ -19,8 +19,7 @@ setup(
    install_requires=[
        "torch>=1.1.0",
        "tqdm",
-        "test-tube>=0.653",
-        "tensorflow>=1.14.0"
+        "test-tube>=0.6.7.4",
    ],
    packages=find_packages(),
    long_description=open("README.md", encoding="utf-8").read(),
@@ -0,0 +1,58 @@
+# Pytorch-Lightning Tests
+
+## Running tests
+The automatic travis tests ONLY run CPU-based tests. Although these cover most of the use cases,
+run on a 2-GPU machine to validate the full test-suite.
+
+
+To run all tests do the following:
+```bash
+git clone https://github.com/williamFalcon/pytorch-lightning
+cd pytorch-lightning
+
+# install module locally
+pip install -e .
+
+# install dev deps
+pip install -r requirements.txt
+
+# run tests
+py.test
+```
+
+To test models that require GPU make sure to run the above command on a GPU machine.
+The GPU machine must have:
+1. At least 2 GPUs.
+2. [NVIDIA-apex](https://github.com/NVIDIA/apex#linux) installed.
+
+
+### test_models.py
+This file fits a tiny model on MNIST using these different set-ups.
+1. CPU only.
+2. Single GPU with DP.
+3. Multiple (2) GPUs using DP.
+3. Multiple (2) GPUs using DDP.
+3. Multiple (2) GPUs using DP + apex (for 16-bit precision).
+3. Multiple (2) GPUs using DDP + apex (for 16-bit precision).   
+
+For each set up it also tests:
+1. model saving.
+2. model loading.
+3. predicting with a loaded model.
+4. simulated save from HPC signal.
+5. simulated load from HPC signal.
+
+## Running Coverage   
+
+```bash
+cd pytorch-lightning
+
+# generate coverage 
+pip install coverage
+coverage run tests/test_models.py   
+
+# print coverage stats
+coverage report -m   
+```
+
+
@@ -0,0 +1,180 @@
+import pytest
+from pytorch_lightning import Trainer
+from pytorch_lightning.examples.new_project_templates.lightning_module_template import LightningTemplateModel
+from argparse import Namespace
+from test_tube import Experiment
+from pytorch_lightning.callbacks import ModelCheckpoint
+import numpy as np
+import warnings
+import torch
+import os
+import shutil
+import pdb
+
+import pytorch_lightning as ptl
+import torch
+from torch.nn import functional as F
+from torch.utils.data import DataLoader
+from torchvision.datasets import MNIST
+
+
+class CoolModel(ptl.LightningModule):
+
+    def __init(self):
+        super(CoolModel, self).__init__()
+        # not the best model...
+        self.l1 = torch.nn.Linear(28 * 28, 10)
+
+    def forward(self, x):
+        return torch.relu(self.l1(x))
+
+    def my_loss(self, y_hat, y):
+        return F.cross_entropy(y_hat, y)
+
+    def training_step(self, batch, batch_nb):
+        x, y = batch
+        y_hat = self.forward(x)
+        return {'tng_loss': self.my_loss(y_hat, y)}
+
+    def validation_step(self, batch, batch_nb):
+        x, y = batch
+        y_hat = self.forward(x)
+        return {'val_loss': self.my_loss(y_hat, y)}
+
+    def validation_end(self, outputs):
+        avg_loss = torch.stack([x for x in outputs['val_loss']]).mean()
+        return avg_loss
+
+    def configure_optimizers(self):
+        return [torch.optim.Adam(self.parameters(), lr=0.02)]
+
+    @ptl.data_loader
+    def tng_dataloader(self):
+        return DataLoader(MNIST('path/to/save', train=True), batch_size=32)
+
+    @ptl.data_loader
+    def val_dataloader(self):
+        return DataLoader(MNIST('path/to/save', train=False), batch_size=32)
+
+    @ptl.data_loader
+    def test_dataloader(self):
+        return DataLoader(MNIST('path/to/save', train=False), batch_size=32)
+
+
+def get_model():
+    # set up model with these hyperparams
+    root_dir = os.path.dirname(os.path.realpath(__file__))
+    hparams = Namespace(**{'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})
+    model = LightningTemplateModel(hparams)
+
+    return model, hparams
+
+
+def get_exp(debug=True):
+    # 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')
+    return exp
+
+
+def init_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)
+
+    os.makedirs(save_dir, exist_ok=True)
+
+    return save_dir
+
+
+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)
+
+
+def load_model(exp, save_dir):
+
+    # load trained model
+    tags_path = exp.get_data_path(exp.name, exp.version)
+    tags_path = os.path.join(tags_path, 'meta_tags.csv')
+
+    checkpoints = [x for x in os.listdir(save_dir) if '.ckpt' in x]
+    weights_dir = os.path.join(save_dir, checkpoints[0])
+
+    trained_model = LightningTemplateModel.load_from_metrics(weights_path=weights_dir, tags_csv=tags_path, on_gpu=True)
+
+    assert trained_model is not None, 'loading model failed'
+
+    return trained_model
+
+
+def run_prediction(dataloader, trained_model):
+    # run prediction on 1 batch
+    for batch in dataloader:
+        break
+
+    x, y = batch
+    x = x.view(x.size(0), -1)
+
+    y_hat = trained_model(x)
+
+    # 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()
+
+    print(val_acc)
+
+    assert val_acc > 0.70, f'this model is expected to get > 0.7 in test set (it got {val_acc})'
+
+
+def main():
+
+    save_dir = init_save_dir()
+
+    # exp file to get meta
+    exp = get_exp(False)
+    exp.save()
+
+    # exp file to get weights
+    checkpoint = ModelCheckpoint(save_dir)
+
+    trainer = Trainer(
+        experiment=exp,
+        checkpoint_callback=checkpoint,
+        progress_bar=True,
+        max_nb_epochs=1,
+        gpus=[0, 1],
+        distributed_backend='dp',
+    )
+
+    model = CoolModel()
+
+    result = trainer.fit(model)
+
+    # correct result and ok accuracy
+    assert result == 1, 'amp + ddp model failed to complete'
+
+    # test model loading
+    pretrained_model = load_model(exp, save_dir)
+
+    # test model preds
+    run_prediction(model.test_dataloader, pretrained_model)
+
+    clear_save_dir()
+
+
+if __name__ == '__main__':
+    main()
@@ -0,0 +1,688 @@
+import pytest
+from pytorch_lightning import Trainer
+from pytorch_lightning.examples.new_project_templates.lightning_module_template import LightningTemplateModel
+from pytorch_lightning.testing_models.lm_test_module import LightningTestModel
+from argparse import Namespace
+from test_tube import Experiment, SlurmCluster
+from pytorch_lightning.callbacks import ModelCheckpoint, EarlyStopping
+from pytorch_lightning.utils.debugging import MisconfigurationException
+from pytorch_lightning.root_module import memory
+from pytorch_lightning.models.trainer import reduce_distributed_output
+from pytorch_lightning.root_module import model_saving
+import numpy as np
+import warnings
+import torch
+import os
+import shutil
+import pdb
+
+SEED = 2334
+torch.manual_seed(SEED)
+np.random.seed(SEED)
+
+
+# ------------------------------------------------------------------------
+# TESTS
+# ------------------------------------------------------------------------
+
+def test_cpu_slurm_save_load():
+    """
+    Verify model save/load/checkpoint on CPU
+    :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()
+
+    cluster_a = SlurmCluster()
+    trainer_options = dict(
+        max_nb_epochs=1,
+        cluster=cluster_a,
+        experiment=exp,
+        checkpoint_callback=ModelCheckpoint(save_dir)
+    )
+
+    # fit model
+    trainer = Trainer(**trainer_options)
+    result = trainer.fit(model)
+    real_global_step = trainer.global_step
+
+    # traning complete
+    assert result == 1, 'amp + ddp model failed to complete'
+
+    # predict with trained model before saving
+    # make a prediction
+    for batch in model.test_dataloader:
+        break
+
+    x, y = batch
+    x = x.view(x.size(0), -1)
+
+    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)
+    trainer_options = dict(
+        max_nb_epochs=1,
+        cluster=SlurmCluster(continue_tng_hparams),
+        experiment=exp,
+        checkpoint_callback=ModelCheckpoint(save_dir),
+    )
+    trainer = Trainer(**trainer_options)
+    model = LightningTestModel(hparams)
+
+    # set the epoch start hook so we can predict before the model does the full training
+    def assert_pred_same():
+        assert trainer.global_step == real_global_step and trainer.global_step > 0
+
+        # predict with loaded model to make sure answers are the same
+        trainer.model.eval()
+        new_pred = trainer.model(x)
+        assert torch.all(torch.eq(pred_before_saving, new_pred)).item() == 1
+
+    model.on_epoch_start = assert_pred_same
+
+    # by calling fit again, we trigger training, loading weights from the cluster
+    # and our hook to predict using current model before any more weight updates
+    trainer.fit(model)
+
+    clear_save_dir()
+
+
+def test_loading_meta_tags():
+    hparams = get_hparams()
+
+    save_dir = init_save_dir()
+
+    # save tags
+    exp = get_exp(False)
+    exp.tag({'some_str':'a_str', 'an_int': 1, 'a_float': 2.0})
+    exp.argparse(hparams)
+    exp.save()
+
+    # 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)
+
+    assert tags.batch_size == 32 and tags.hidden_dim == 1000
+
+    clear_save_dir()
+
+
+def test_dp_output_reduce():
+
+    # test identity when we have a single gpu
+    out = torch.rand(3, 1)
+    assert reduce_distributed_output(out, nb_gpus=1) is out
+
+    # average when we have multiples
+    assert reduce_distributed_output(out, nb_gpus=2) == out.mean()
+
+    # when we have a dict of vals
+    out = {
+        'a': out,
+        'b': {
+            'c': out
+        }
+    }
+    reduced = reduce_distributed_output(out, nb_gpus=3)
+    assert reduced['a'] == out['a']
+    assert reduced['b']['c'] == out['b']['c']
+
+
+def test_model_saving_loading():
+    """
+    Tests use case where trainer saves the model, and user loads it from tags independently
+    :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()
+
+    trainer_options = dict(
+        max_nb_epochs=1,
+        cluster=SlurmCluster(),
+        experiment=exp,
+        checkpoint_callback=ModelCheckpoint(save_dir)
+    )
+
+    # fit model
+    trainer = Trainer(**trainer_options)
+    result = trainer.fit(model)
+
+    # traning complete
+    assert result == 1, 'amp + ddp model failed to complete'
+
+    # make a prediction
+    for batch in model.test_dataloader:
+        break
+
+    x, y = batch
+    x = x.view(x.size(0), -1)
+
+    # generate preds before saving model
+    model.eval()
+    pred_before_saving = model(x)
+
+    # save model
+    new_weights_path = os.path.join(save_dir, 'save_test.ckpt')
+    trainer.save_checkpoint(new_weights_path)
+
+    # load new model
+    tags_path = exp.get_data_path(exp.name, exp.version)
+    tags_path = os.path.join(tags_path, 'meta_tags.csv')
+    model_2 = LightningTestModel.load_from_metrics(weights_path=new_weights_path, tags_csv=tags_path, on_gpu=False)
+    model_2.eval()
+
+    # make prediction
+    # assert that both predictions are the same
+    new_pred = model_2(x)
+    assert torch.all(torch.eq(pred_before_saving, new_pred)).item() == 1
+
+    clear_save_dir()
+
+
+
+
+def test_model_freeze_unfreeze():
+    hparams = get_hparams()
+    model = LightningTestModel(hparams)
+
+    model.freeze()
+    model.unfreeze()
+
+
+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')
+        return
+
+    # simulate setting slurm flags
+    os.environ['MASTER_PORT'] = str(np.random.randint(12000, 19000, 1)[0])
+    os.environ['SLURM_LOCALID'] = str(0)
+
+    hparams = get_hparams()
+    model = LightningTestModel(hparams)
+
+    trainer_options = dict(
+        progress_bar=True,
+        max_nb_epochs=1,
+        gpus=[0],
+        distributed_backend='ddp',
+        use_amp=True
+    )
+
+    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['checkpoint_callback'] = checkpoint
+    trainer_options['experiment'] = exp
+
+    # fit model
+    trainer = Trainer(**trainer_options)
+    trainer.is_slurm_managing_tasks = True
+    result = trainer.fit(model)
+
+    # correct result and ok accuracy
+    assert result == 1, 'amp + ddp model failed to complete'
+
+    # test root model address
+    assert trainer.resolve_root_node_address('abc') == 'abc'
+    assert trainer.resolve_root_node_address('abc[23]') == 'abc23'
+    assert trainer.resolve_root_node_address('abc[23-24]') == 'abc23'
+    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)
+
+    # test model preds
+    run_prediction(model.test_dataloader, pretrained_model)
+
+    if trainer.use_ddp:
+        # on hpc this would work fine... but need to hack it for the purpose of the test
+        trainer.model = pretrained_model
+        trainer.optimizers = pretrained_model.configure_optimizers()
+
+    # test HPC loading / saving
+    trainer.hpc_save(save_dir, exp)
+    trainer.hpc_load(save_dir, on_gpu=True)
+
+    # test freeze on gpu
+    model.freeze()
+    model.unfreeze()
+
+    clear_save_dir()
+
+
+def test_early_stopping_cpu_model():
+    """
+    Test each of the trainer options
+    :return:
+    """
+
+    stopping = EarlyStopping()
+    trainer_options = dict(
+        early_stop_callback=stopping,
+        gradient_clip=1.0,
+        overfit_pct=0.20,
+        track_grad_norm=2,
+        print_nan_grads=True,
+        progress_bar=False,
+        experiment=get_exp(),
+        train_percent_check=0.1,
+        val_percent_check=0.1
+    )
+
+    model, hparams = get_model()
+    run_gpu_model_test(trainer_options, model, hparams, on_gpu=False)
+
+    # test freeze on cpu
+    model.freeze()
+    model.unfreeze()
+
+
+def test_cpu_model_with_amp():
+    """
+    Make sure model trains on CPU
+    :return:
+    """
+
+    trainer_options = dict(
+        progress_bar=False,
+        experiment=get_exp(),
+        max_nb_epochs=1,
+        train_percent_check=0.4,
+        val_percent_check=0.4,
+        use_amp=True
+    )
+
+    model, hparams = get_model()
+
+    with pytest.raises((MisconfigurationException, ModuleNotFoundError)):
+        run_gpu_model_test(trainer_options, model, hparams, on_gpu=False)
+
+
+def test_cpu_model():
+    """
+    Make sure model trains on CPU
+    :return:
+    """
+
+    trainer_options = dict(
+        progress_bar=False,
+        experiment=get_exp(),
+        max_nb_epochs=1,
+        train_percent_check=0.4,
+        val_percent_check=0.4
+    )
+
+    model, hparams = get_model()
+
+    run_gpu_model_test(trainer_options, model, hparams, on_gpu=False)
+
+
+def test_all_features_cpu_model():
+    """
+    Test each of the trainer options
+    :return:
+    """
+
+    trainer_options = dict(
+        gradient_clip=1.0,
+        overfit_pct=0.20,
+        track_grad_norm=2,
+        print_nan_grads=True,
+        progress_bar=False,
+        experiment=get_exp(),
+        max_nb_epochs=1,
+        train_percent_check=0.4,
+        val_percent_check=0.4
+    )
+
+    model, hparams = get_model()
+    run_gpu_model_test(trainer_options, model, hparams, on_gpu=False)
+
+
+def test_single_gpu_model():
+    """
+    Make sure single GPU works (DP mode)
+    :return:
+    """
+    if not torch.cuda.is_available():
+        warnings.warn('test_single_gpu_model cannot run. Rerun on a GPU node to run this test')
+        return
+    model, hparams = get_model()
+
+    trainer_options = dict(
+        progress_bar=False,
+        max_nb_epochs=1,
+        train_percent_check=0.1,
+        val_percent_check=0.1,
+        gpus=[0]
+    )
+
+    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')
+        return
+    model, hparams = get_model()
+    trainer_options = dict(
+        progress_bar=False,
+        max_nb_epochs=1,
+        train_percent_check=0.1,
+        val_percent_check=0.1,
+        gpus='-1'
+    )
+
+    run_gpu_model_test(trainer_options, model, hparams)
+
+    # test memory helper functions
+    memory.get_gpu_memory_map()
+
+
+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')
+        return
+    model, hparams = get_model()
+    trainer_options = dict(
+        max_nb_epochs=1,
+        gpus='0, 1',  # test init with gpu string
+        distributed_backend='dp',
+        use_amp=True
+    )
+    with pytest.raises(MisconfigurationException):
+        run_gpu_model_test(trainer_options, model, hparams)
+
+
+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')
+        return
+
+    os.environ['MASTER_PORT'] = str(np.random.randint(12000, 19000, 1)[0])
+    model, hparams = get_model()
+    trainer_options = dict(
+        progress_bar=False,
+        max_nb_epochs=1,
+        train_percent_check=0.4,
+        val_percent_check=0.2,
+        gpus=[0, 1],
+        distributed_backend='ddp'
+    )
+
+    run_gpu_model_test(trainer_options, model, hparams)
+
+
+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')
+        return
+
+    os.environ['MASTER_PORT'] = str(np.random.randint(12000, 19000, 1)[0])
+
+    hparams = get_hparams()
+    model = LightningTestModel(hparams)
+
+    trainer_options = dict(
+        progress_bar=True,
+        max_nb_epochs=1,
+        gpus=[0, 1],
+        distributed_backend='ddp',
+        use_amp=True
+    )
+
+    run_gpu_model_test(trainer_options, model, hparams)
+
+
+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')
+        return
+
+    os.environ['MASTER_PORT'] = str(np.random.randint(12000, 19000, 1)[0])
+
+    hparams = get_hparams()
+    model = LightningTestModel(hparams, force_remove_distributed_sampler=True)
+
+    exp = get_exp(True)
+    exp.save()
+
+    trainer = Trainer(
+        experiment=exp,
+        progress_bar=False,
+        max_nb_epochs=1,
+        gpus=[0, 1],
+        distributed_backend='ddp',
+        use_amp=True
+    )
+
+    with pytest.raises(MisconfigurationException):
+        trainer.get_dataloaders(model)
+
+    clear_save_dir()
+
+
+# ------------------------------------------------------------------------
+# UTILS
+# ------------------------------------------------------------------------
+def run_gpu_model_test(trainer_options, model, hparams, on_gpu=True):
+    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['checkpoint_callback'] = checkpoint
+    trainer_options['experiment'] = exp
+
+    # fit model
+    trainer = Trainer(**trainer_options)
+    result = trainer.fit(model)
+
+    # correct result and ok accuracy
+    assert result == 1, 'amp + ddp model failed to complete'
+
+    # test model loading
+    pretrained_model = load_model(exp, save_dir, on_gpu)
+
+    # test model preds
+    run_prediction(model.test_dataloader, pretrained_model)
+
+    if trainer.use_ddp:
+        # on hpc this would work fine... but need to hack it for the purpose of the test
+        trainer.model = pretrained_model
+        trainer.optimizers = pretrained_model.configure_optimizers()
+
+    # test HPC loading / saving
+    trainer.hpc_save(save_dir, exp)
+    trainer.hpc_load(save_dir, on_gpu=on_gpu)
+
+    clear_save_dir()
+
+
+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 get_model():
+    # set up model with these hyperparams
+    hparams = get_hparams()
+    model = LightningTemplateModel(hparams)
+
+    return model, hparams
+
+
+def get_exp(debug=True):
+    # 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')
+    return exp
+
+
+def init_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)
+
+    os.makedirs(save_dir, exist_ok=True)
+
+    return save_dir
+
+
+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)
+
+
+def load_model(exp, save_dir, on_gpu, map_location=None):
+
+    # load trained model
+    tags_path = exp.get_data_path(exp.name, exp.version)
+    tags_path = os.path.join(tags_path, 'meta_tags.csv')
+
+    checkpoints = [x for x in os.listdir(save_dir) if '.ckpt' in x]
+    weights_dir = os.path.join(save_dir, checkpoints[0])
+
+    trained_model = LightningTemplateModel.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'
+
+    return trained_model
+
+
+def run_prediction(dataloader, trained_model):
+    # run prediction on 1 batch
+    for batch in dataloader:
+        break
+
+    x, y = batch
+    x = x.view(x.size(0), -1)
+
+    y_hat = trained_model(x)
+
+    # 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()
+
+    print(val_acc)
+
+    assert val_acc > 0.50, f'this model is expected to get > 0.50 in test set (it got {val_acc})'
+
+
+def assert_ok_acc(trainer):
+    # this model should get 0.80+ acc
+    acc = trainer.tng_tqdm_dic['val_acc']
+    assert acc > 0.50, f'model failed to get expected 0.50 validation accuracy. Got: {acc}'
+
+
+if __name__ == '__main__':
+    pytest.main([__file__])
				`@@ -1 +0,0 @@`
				`from .lightning_module_template import LightningTemplateModel`
				`@@ -0,0 +1 @@`
				`from .new_project_templates.lightning_module_template import LightningTemplateModel`