release v0.112

.gitignore

@@ -116,3 +116,6 @@ ENV/
 
 # mypy
 .mypy_cache/
+
+# data
+mnist/

README.md

@@ -20,221 +20,151 @@ pip install pytorch-lightning
 ```
 
 ## Docs   
-In progress. Documenting now!  
+**[View the docs here](https://williamfalcon.github.io/pytorch-lightning/)**
-
-## Disclaimer
-This is a research tool I built for myself internally while doing my PhD. The API is not 100% production quality, but my hope is that by open-sourcing, we can all get it there (I don't have too much time nowadays to write production-level code).
 
 ## What is it?  
-Keras is too abstract for researchers. Lightning makes it so you only have to define your model but still control all details of training if you need to.   
+Keras and fast.ai are too abstract for researchers. Lightning abstracts the full training loop but gives you control in the critical points.   
 
-Pytorch   
-<-- Lightning   
-Your model.   
 
-**Lightning will do the following for you:**   
+## 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.
 
-1. Run the training loop.   
-2. Run the validation loop.   
-3. Run the testing loop.   
-4. Early stopping.   
-5. Learning rate annealing.
-6. Can train complex models like GANs or anything with multiple optimizers.
-7. Weight checkpointing.   
-8. Model saving.   
-9. Model loading.   
-10. Log training details (through test-tube).  
-11. Run training on multiple GPUs (through test-tube).   
-12. Run training on a GPU cluster managed by SLURM (through test-tube).   
-13. Distribute memory-bound models on multiple GPUs.
-14. Give your model hyperparameters parsed from the command line OR a JSON file.   
-15. Run your model in a dev environment where nothing logs.      
-
-## Usage
 To use lightning do 2 things:  
-1. [Define a trainer](https://github.com/williamFalcon/pytorch-lightning/blob/master/docs/source/examples/basic_trainer.py) (which will run ALL your models).   
+1. [Define a Trainer](https://github.com/williamFalcon/pytorch-lightning/blob/master/examples/new_project_templates/trainer_cpu_template.py).   
-2. [Define a model](https://github.com/williamFalcon/pytorch-lightning/blob/master/docs/source/examples/example_model.py).     
+2. [Define a LightningModel](https://github.com/williamFalcon/pytorch-lightning/blob/master/examples/new_project_templates/lightning_module_template.py).     
 
-#### Basic trainer example  
+## What does lightning control for me?
-See [this demo](https://github.com/williamFalcon/pytorch-lightning/blob/master/docs/source/examples/fully_featured_trainer.py) for a more robust trainer example.   
+Everything! Except the following three things:
+
+**What happens in the training loop**    
 
 ```python
-import os
+# define what happens for training here
-import sys
+def training_step(self, data_batch, batch_nb):
+    x, y = data_batch
     
-from test_tube import HyperOptArgumentParser, Experiment
+    # define your own forward and loss calculation
-from pytorch_lightning.models.trainer import Trainer
+    out = self.forward(x)
-from pytorch_lightning.utils.arg_parse import add_default_args
+    loss = my_loss(out, y)
-from pytorch_lightning.utils.pt_callbacks import EarlyStopping, ModelCheckpoint
+    return {'loss': loss} 
-from demo.example_model import ExampleModel
+```
 
+**What happens in the validation loop**      
 
-def main(hparams):
+```python
+# define what happens for validation here
+def validation_step(self, data_batch, batch_nb):    
+    x, y = data_batch
+    
+    # define your own forward and loss calculation
+    out = self.forward(x)
+    loss = my_loss(out, y)
+    return {'loss': loss} 
+```
+
+**And what to do with the output of all validation batches**    
+
+```python
+def validation_end(self, outputs):
     """
-    Main training routine specific for this project
+    Called at the end of validation to aggregate outputs
-    :param hparams:
+    :param outputs: list of individual outputs of each validation step
     :return:
     """
-    # init experiment
+    val_loss_mean = 0
-    exp = Experiment(
+    val_acc_mean = 0
-        name=hparams.tt_name,
-        debug=hparams.debug,
-        save_dir=hparams.tt_save_path,
-        version=hparams.hpc_exp_number,
-        autosave=False,
-        description=hparams.tt_description
-    )
-
-    exp.argparse(hparams)
-    exp.save()
-
-    model_save_path = '{}/{}/{}'.format(hparams.model_save_path, exp.name, exp.version)
-
-    # build model
-    model = ExampleModel(hparams)
-
-    # callbacks
-    early_stop = EarlyStopping(monitor='val_acc', patience=3, mode='min', verbose=True)
-    checkpoint = ModelCheckpoint(filepath=model_save_path, save_function=None, save_best_only=True, verbose=True, monitor='val_acc', mode='min')
-
-    # configure trainer
-    trainer = Trainer(experiment=exp, checkpoint_callback=checkpoint, early_stop_callback=early_stop)
-
-    # train model
-    trainer.fit(model)
-
-
-if __name__ == '__main__':
-
-    # use default args given by lightning
-    root_dir = os.path.split(os.path.dirname(sys.modules['__main__'].__file__))[0]
-    parent_parser = HyperOptArgumentParser(strategy='random_search', add_help=False)
-    add_default_args(parent_parser, root_dir)
-
-    # allow model to overwrite or extend args
-    parser = ExampleModel.add_model_specific_args(parent_parser)
-    hyperparams = parser.parse_args()
-
-    # train model
-    main(hyperparams)
-
-```
-
-#### Basic model example  
-Here we only show the method signatures. It's up to you to define the content.   
-
-```python
-from torch import nn
-
-class My_Model(RootModule):
-    def __init__(self):
-        # define model
-        self.l1 = nn.Linear(200, 10)
-
-    # ---------------
-    # TRAINING
-    def training_step(self, data_batch):
-        x, y = data_batch
-        y_hat = self.l1(x)
-        loss = some_loss(y_hat)
-
-        return loss_val, {'train_loss': loss}
-
-    def validation_step(self, data_batch):
-        x, y = data_batch
-        y_hat = self.l1(x)
-        loss = some_loss(y_hat)
-
-        return loss_val, {'val_loss': loss}
-
-     def validation_end(self, outputs):
-        total_accs = []
-
     for output in outputs:
-            total_accs.append(output['val_acc'].item())
+        val_loss_mean += output['val_loss']
+        val_acc_mean += output['val_acc']
 
-        # return a dict
+    val_loss_mean /= len(outputs)
-        return {'total_acc': np.mean(total_accs)}
+    val_acc_mean /= len(outputs)
+    tqdm_dic = {'val_loss': val_loss_mean.item(), 'val_acc': val_acc_mean.item()}
+    return tqdm_dic
+```
 
-     # ---------------
+## Lightning gives you options to control the following:
-     # SAVING
-     def get_save_dict(self):
-        # lightning saves for you. Here's your chance to say what you want to save
-        checkpoint = {'state_dict': self.state_dict()}
 
-        return checkpoint
+###### Checkpointing    
 
-    def load_model_specific(self, checkpoint):
+- [Model saving](https://williamfalcon.github.io/pytorch-lightning/Trainer/Checkpointing/#model-saving)
-        # lightning loads for you. Here's your chance to say what you want to load
+- [Model loading](https://williamfalcon.github.io/pytorch-lightning/LightningModule/methods/#load-from-metrics) 
-        self.load_state_dict(checkpoint['state_dict'])
 
-    # ---------------
+###### Computing cluster (SLURM)    
-    # TRAINING CONFIG
-    def configure_optimizers(self):
-        # give lightning the list of optimizers you want to use.
-        # lightning will call automatically
-        optimizer = self.choose_optimizer('adam', self.parameters(), {'lr': self.hparams.learning_rate}, 'optimizer')
-        return [optimizer]
 
-    @property
+- [Running grid search on a cluster](https://williamfalcon.github.io/pytorch-lightning/Trainer/SLURM%20Managed%20Cluster#running-grid-search-on-a-cluster) 
-    def tng_dataloader(self):
+- [Walltime auto-resubmit](https://williamfalcon.github.io/pytorch-lightning/Trainer/SLURM%20Managed%20Cluster#walltime-auto-resubmit)   
-        return pytorch_dataloader('train')
 
-    @property
+###### Debugging  
-    def val_dataloader(self):
-        return pytorch_dataloader('val')
 
-    @property
+- [Fast dev run](https://williamfalcon.github.io/pytorch-lightning/Trainer/debugging/#fast-dev-run)
-    def test_dataloader(self):
+- [Inspect gradient norms](https://williamfalcon.github.io/pytorch-lightning/Trainer/debugging/#inspect-gradient-norms)
-        return pytorch_dataloader('test')
+- [Log GPU usage](https://williamfalcon.github.io/pytorch-lightning/Trainer/debugging/#Log-gpu-usage)
+- [Make model overfit on subset of data](https://williamfalcon.github.io/pytorch-lightning/Trainer/debugging/#make-model-overfit-on-subset-of-data)
+- [Print the parameter count by layer](https://williamfalcon.github.io/pytorch-lightning/Trainer/debugging/#print-the-parameter-count-by-layer)
+- [Pring which gradients are nan](https://williamfalcon.github.io/pytorch-lightning/Trainer/debugging/#print-which-gradients-are-nan)
 
-    # ---------------
+
-    # MODIFY YOUR COMMAND LINE ARGS
+###### Distributed training    
-    @staticmethod
+
-    def add_model_specific_args(parent_parser):    
+- [16-bit mixed precision](https://williamfalcon.github.io/pytorch-lightning/Trainer/Distributed%20training/#16-bit-mixed-precision)
-        parser = HyperOptArgumentParser(strategy=parent_parser.strategy, parents=[parent_parser])
+- [Multi-GPU](https://williamfalcon.github.io/pytorch-lightning/Trainer/Distributed%20training/#Multi-GPU)
-        parser.add_argument('--out_features', default=20)
+- [Multi-node](https://williamfalcon.github.io/pytorch-lightning/Trainer/Distributed%20training/#Multi-node)
-        return parser
+- [Single GPU](https://williamfalcon.github.io/pytorch-lightning/Trainer/Distributed%20training/#single-gpu)
+- [Self-balancing architecture](https://williamfalcon.github.io/pytorch-lightning/Trainer/Distributed%20training/#self-balancing-architecture)
+
+
+###### Experiment Logging   
+
+- [Display metrics in progress bar](https://williamfalcon.github.io/pytorch-lightning/Trainer/Logging/#display-metrics-in-progress-bar)
+- Log arbitrary metrics
+- [Log metric row every k batches](https://williamfalcon.github.io/pytorch-lightning/Trainer/Logging/#log-metric-row-every-k-batches)
+- [Process position](https://williamfalcon.github.io/pytorch-lightning/Trainer/Logging/#process-position)
+- [Save a snapshot of all hyperparameters](https://williamfalcon.github.io/pytorch-lightning/Trainer/Logging/#save-a-snapshot-of-all-hyperparameters) 
+- [Snapshot code for a training run](https://williamfalcon.github.io/pytorch-lightning/Trainer/Logging/#snapshot-code-for-a-training-run) 
+- [Write logs file to csv every k batches](https://williamfalcon.github.io/pytorch-lightning/Trainer/Logging/#write-logs-file-to-csv-every-k-batches)
+
+###### Training loop    
+
+- [Accumulate gradients](https://williamfalcon.github.io/pytorch-lightning/Trainer/Training%20Loop/#accumulated-gradients)
+- [Anneal Learning rate](https://williamfalcon.github.io/pytorch-lightning/Trainer/Training%20Loop/#anneal-learning-rate)
+- [Force training for min or max epochs](https://williamfalcon.github.io/pytorch-lightning/Trainer/Training%20Loop/#force-training-for-min-or-max-epochs)
+- [Force disable early stop](https://williamfalcon.github.io/pytorch-lightning/Trainer/Training%20Loop/#force-disable-early-stop)
+- [Gradient Clipping](https://williamfalcon.github.io/pytorch-lightning/Trainer/Training%20Loop/#gradient-clipping)
+- [Use multiple optimizers (like GANs)](https://williamfalcon.github.io/pytorch-lightning/Pytorch-Lightning/LightningModule/#configure_optimizers)
+- [Set how much of the training set to check (1-100%)](https://williamfalcon.github.io/pytorch-lightning/Trainer/Training%20Loop/#set-how-much-of-the-training-set-to-check)
+
+###### Validation loop    
+
+- [Check validation every n epochs](https://williamfalcon.github.io/pytorch-lightning/Trainer/Validation%20loop/#check-validation-every-n-epochs)
+- [Set how much of the validation set to check](https://williamfalcon.github.io/pytorch-lightning/Trainer/Validation%20loop/#set-how-much-of-the-validation-set-to-check)
+- [Set how much of the test set to check](https://williamfalcon.github.io/pytorch-lightning/Trainer/Validation%20loop/#set-how-much-of-the-test-set-to-check)
+- [Set validation check frequency within 1 training epoch](https://williamfalcon.github.io/pytorch-lightning/Trainer/Validation%20loop/#set-validation-check-frequency-within-1-training-epoch)
+- [Set the number of validation sanity steps](https://williamfalcon.github.io/pytorch-lightning/Trainer/Validation%20loop/#set-the-number-of-validation-sanity-steps)
+
+
+## Demo
+```bash
+# install lightning
+pip install pytorch-lightning
+
+# clone lightning for the demo
+git clone https://github.com/williamFalcon/pytorch-lightning.git
+cd examples/new_project_templates/
+
+# run demo (on cpu)
+python trainer_gpu_cluster_template.py
+```
+
+Without changing the model AT ALL, you can run the model on a single gpu, over multiple gpus, or over multiple nodes.
+```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
 ```
 
 
-### Details    
-
-#### Model definition   
-| Name  | Description  |  Input |  Return |
-|---|---|---|---|
-| training_step  |  Called with a batch of data during training | data from your dataloaders  | tuple: scalar, dict  |
-| validation_step  |  Called with a batch of data during validation  | data from your dataloaders  | tuple: scalar, dict |
-| validation_end  |  Collate metrics from all validation steps |  outputs: array where each item is the output of a validation step | dict: for logging |
-| get_save_dict  | called when your model needs to be saved (checkpoints, hpc save, etc...)  | None  |  dict to be saved |
-
-#### Model training   
-| Name  | Description  |  Input |  Return |
-|---|---|---|---|
-| configure_optimizers  |  called during training setup | None | list: optimizers you want to use |
-| tng_dataloader  |  called during training  | None  | pytorch dataloader |
-| val_dataloader  |  called during validation  | None  | pytorch dataloader |
-| test_dataloader  |  called during testing  | None  | pytorch dataloader |    
-| add_model_specific_args  |  called with args you defined in your main. This lets you tailor args for each model and keep main the same  | argparse  | argparse |
-
-#### Model Saving/Loading   
-| Name  | Description  |  Input |  Return |
-|---|---|---|---|
-| get_save_dict  | called when your model needs to be saved (checkpoints, hpc save, etc...)  | None  |  dict to be saved |
-|  load_model_specific |  called when loading a model | checkpoint: dict you created in get_save_dict  | dict: modified in whatever way you want  |
-
-## Optional model hooks.   
-Add these to the model whenever you want to configure training behavior.    
-
-
-### Model lifecycle hooks
-Use these hooks to customize functionality
-
-| Method | Purpose  | Input  | Output  | Required  |
-|---|---|---|---|---|
-| on_batch_start()  | called right before the batch starts | - | -  | N  |
-| on_batch_end()  | called right after the batch ends | - | -  | N  |
-| on_epoch_start()  | called right before the epoch starts | - | -  | N  |
-| on_epoch_end()  | called right afger the epoch ends | - | -  | N  |
-| on_pre_performance_check()  | called right before the performance check starts | - | -  | N  |
-| on_post_performance_check()  | called right after the batch starts | - | -  | N  |

docs/LightningModule/RequiredTrainerInterface.md

@@ -0,0 +1,394 @@
+# Lightning Module interface
+[[Github Code](https://github.com/williamFalcon/pytorch-lightning/blob/master/pytorch_lightning/root_module/root_module.py)]
+
+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. 
+
+Otherwise, to Define a Lightning Module, implement the following methods:
+
+**Required**:  
+
+- [training_step](RequiredTrainerInterface.md#training_step)   
+- [validation_step](RequiredTrainerInterface.md#validation_step)
+- [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)
+- [test_dataloader](RequiredTrainerInterface.md#test_dataloader)
+
+**Optional**:   
+
+- [update_tng_log_metrics](RequiredTrainerInterface.md#update_tng_log_metrics)
+- [add_model_specific_args](RequiredTrainerInterface.md#add_model_specific_args)
+
+---
+
+### training_step
+
+``` {.python}
+def training_step(self, data_batch, batch_nb)
+```
+
+In this step you'd normally do the forward pass and calculate the loss for a batch. You can also do fancier things like multiple forward passes or something specific to your model.
+
+**Params**    
+
+| Param  | description  |
+|---|---|
+|  data_batch | The output of your dataloader. A tensor, tuple or list  |
+|  batch_nb | Integer displaying which batch this is  |
+
+**Return**   
+
+Dictionary or OrderedDict   
+
+| key  | value  | is required |
+|---|---|---|
+|  loss | tensor scalar  | Y |
+|  prog | Dict for progress bar display. Must have only tensors  | N |
+
+
+**Example**
+
+``` {.python}
+def training_step(self, data_batch, batch_nb):
+    x, y, z = data_batch
+    
+    # implement your own
+    out = self.forward(x)
+    loss = self.loss(out, x)
+    
+    output = {
+        'loss': loss, # required
+        'prog': {'tng_loss': loss, 'batch_nb': batch_nb} # optional
+    }
+    
+    # return a dict
+    return output
+```
+
+---
+
+### validation_step
+
+``` {.python}
+def validation_step(self, data_batch, batch_nb)
+```
+
+In this step you'd normally do the forward pass and calculate the loss for a batch. You can also do fancier things like multiple forward passes or something specific to your model.
+This is most likely the same as your training_step. But unlike training step, the outputs from here will go to validation_end for collation.
+
+**Params**   
+
+| Param  | description  |
+|---|---|
+|  data_batch | The output of your dataloader. A tensor, tuple or list  |
+|  batch_nb | Integer displaying which batch this is  |
+
+**Return**   
+
+| Return  | description  | optional |
+|---|---|---|   
+|  dict | Dict of OrderedDict with metrics to display in progress bar. All keys must be tensors. | Y |
+
+**Example**
+
+``` {.python}
+def validation_step(self, data_batch, batch_nb):
+    x, y, z = data_batch
+    
+    # implement your own
+    out = self.forward(x)
+    loss = self.loss(out, x)
+    
+    # calculate acc
+    labels_hat = torch.argmax(out, dim=1)
+    val_acc = torch.sum(y == labels_hat).item() / (len(y) * 1.0)
+    
+    # all optional...
+    # return whatever you need for the collation function validation_end
+    output = OrderedDict({
+        'val_loss': loss_val,
+        'val_acc': torch.tensor(val_acc), # everything must be a tensor
+    })
+    
+    # return an optional dict
+    return output
+```
+
+--- 
+### validation_end
+
+``` {.python}
+def validation_end(self, outputs)
+```
+
+Called at the end of the validation loop with the output of each validation_step.
+
+**Params**    
+
+| Param  | description  |
+|---|---|
+|  outputs | List of outputs you defined in validation_step |
+
+**Return**   
+
+| Return  | description  | optional |
+|---|---|---|   
+|  dict | Dict of OrderedDict with metrics to display in progress bar | Y |
+
+**Example**
+
+``` {.python}
+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:
+    """
+    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
+```
+
+--- 
+### configure_optimizers 
+
+``` {.python}
+def configure_optimizers(self)
+```
+
+Set up as many optimizers as you need. Normally you'd need one. But in the case of GANs or something more esoteric you might have multiple. 
+Lightning will call .backward() and .step() on each one.  If you use 16 bit precision it will also handle that.
+
+
+##### Return
+List - List of optimizers
+
+**Example**
+
+``` {.python}
+# most cases
+def configure_optimizers(self):
+    opt = Adam(lr=0.01)
+    return [opt]
+    
+# gan example
+def configure_optimizers(self):
+    generator_opt = Adam(lr=0.01)
+    disriminator_opt = Adam(lr=0.02)
+    return [generator_opt, disriminator_opt] 
+```
+
+--- 
+### get_save_dict 
+
+``` {.python}
+def get_save_dict(self)
+```
+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.
+
+##### Return
+Dictionary - No required keys. Most of the time as described in this example.   
+
+**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
+```
+
+--- 
+### load_model_specific 
+
+``` {.python}
+def load_model_specific(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. 
+
+##### Return
+Nothing  
+
+**Example**
+
+``` {.python}
+def load_model_specific(self, checkpoint):
+    # you defined 'state_dict' in get_save_dict()
+    self.load_state_dict(checkpoint['state_dict'])
+```
+
+--- 
+### tng_dataloader 
+
+``` {.python}
+@property
+def tng_dataloader(self)
+```
+Called by lightning during training loop. Define it as a property.
+
+##### Return
+Pytorch DataLoader
+
+**Example**
+
+``` {.python}
+@property
+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
+```
+
+--- 
+### val_dataloader 
+
+``` {.python}
+@property
+def tng_dataloader(self)
+```
+Called by lightning during validation loop. Define it as a property.
+
+##### Return
+Pytorch DataLoader
+
+**Example**
+
+``` {.python}
+@property
+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
+```
+
+--- 
+### test_dataloader 
+
+``` {.python}
+@property
+def test_dataloader(self)
+```
+Called by lightning during test loop. Define it as a property.
+
+##### Return
+Pytorch DataLoader
+
+**Example**
+
+``` {.python}
+@property
+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
+```
+
+--- 
+### update_tng_log_metrics 
+
+``` {.python}
+def update_tng_log_metrics(self, logs)
+```
+Called by lightning right before it logs metrics for this batch.
+This is a chance to ammend or add to the metrics about to be logged.
+
+##### Return
+Dict 
+
+**Example**
+
+``` {.python}
+def update_tng_log_metrics(self, logs):
+    # modify or add to logs
+    return logs
+```
+
+--- 
+### add_model_specific_args 
+
+``` {.python}
+@staticmethod
+def add_model_specific_args(parent_parser, root_dir)
+```
+Lightning has a list of default argparse commands.
+This method is your chance to add or modify commands specific to your model.
+The [hyperparameter argument parser](https://williamfalcon.github.io/test-tube/hyperparameter_optimization/HyperOptArgumentParser/) is available anywhere in your model by calling self.hparams.
+
+##### Return
+An argument parser
+
+**Example**
+
+``` {.python}
+@staticmethod
+def add_model_specific_args(parent_parser, root_dir):
+    parser = HyperOptArgumentParser(strategy=parent_parser.strategy, parents=[parent_parser])
+
+    # param overwrites
+    # parser.set_defaults(gradient_clip=5.0)
+
+    # 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)
+    parser.add_argument('--out_features', default=10)
+    parser.add_argument('--hidden_dim', default=50000) # 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, type=float, options=[0.0001, 0.0005, 0.001, 0.005],
+                    tunable=False)
+    parser.opt_list('--batch_size', default=256, type=int, options=[32, 64, 128, 256], tunable=False)
+    parser.opt_list('--optimizer_name', default='adam', type=str, options=['adam'], tunable=False)
+    return parser
+```

docs/LightningModule/methods.md

@@ -0,0 +1,49 @@
+Lightning modules are strict superclasses of torch.nn.Module. A LightningModule offers the following in addition to that API.
+
+---    
+### freeze 
+Freeze all params for inference
+```{.python}
+model = MyLightningModule(...)
+model.freeze()
+```
+
+---    
+### load_from_metrics
+This is the easiest/fastest way which uses the meta_tags.csv file from test-tube to rebuild the model.
+The meta_tags.csv file can be found in the test-tube experiment save_dir.       
+
+```{.python}
+pretrained_model = MyLightningModule.load_from_metrics(
+    weights_path='/path/to/pytorch_checkpoint.ckpt',
+    tags_csv='/path/to/test_tube/experiment/version/meta_tags.csv',
+    on_gpu=True,
+    map_location=None
+)
+    
+# predict 
+pretrained_model.freeze()
+y_hat = pretrained_model(x)
+```
+
+**Params**    
+
+| Param  | description  |
+|---|---|
+|  weights_path | Path to a pytorch checkpoint  |
+|  tags_csv | Path to meta_tags.csv file generated by the test-tube Experiment  |
+|  on_gpu | if True, puts model on GPU. Make sure to use transforms option if model devices have changed  |
+|  map_location | A dictionary mapping saved weight GPU devices to new GPU devices |
+
+**Returns**    
+
+LightningModule - The pretrained LightningModule
+
+---    
+### unfreeze 
+Unfreeze all params for inference
+```{.python}
+model = MyLightningModule(...)
+model.unfreeze()
+```
+

docs/LightningModule/properties.md

@@ -0,0 +1,25 @@
+A LightningModule has the following properties which you can access at any time
+
+--- 
+#### current_epoch
+The current epoch   
+
+---
+#### dtype    
+Current dtype    
+
+--- 
+#### global_step
+Total training batches seen across all epochs   
+
+--- 
+#### gradient_clip
+The current gradient clip value    
+
+---
+#### on_gpu    
+True if your model is currently running on GPUs. Useful to set flags around the LightningModule for different CPU vs GPU behavior.    
+
+---
+#### Trainer
+Last resort access to any state the trainer has. Changing certain properties here could affect your training run.

docs/Trainer/Checkpointing.md

@@ -0,0 +1,22 @@
+Lightning can automate saving and loading checkpoints.
+
+---
+### Model saving
+To enable checkpointing, define the checkpoint callback and give it to the trainer.
+
+``` {.python}
+from pytorch_lightning.utils.pt_callbacks import ModelCheckpoint
+
+checkpoint_callback = ModelCheckpoint(
+    filepath='/path/to/store/weights.ckpt',
+    save_best_only=True,
+    verbose=True,
+    monitor='val_loss',
+    mode='min'
+)
+
+trainer = Trainer(checkpoint_callback=checkpoint_callback)
+```
+
+
+

docs/Trainer/Distributed training.md

@@ -0,0 +1,55 @@
+Lightning makes multi-gpu training and 16 bit training trivial.
+
+*Note:*   
+None of the flags below require changing anything about your lightningModel definition. 
+
+---
+#### 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.    
+First, install apex (if install fails, look [here](https://github.com/NVIDIA/apex)):
+```bash
+$ git clone https://github.com/NVIDIA/apex
+$ cd apex
+$ pip install -v --no-cache-dir --global-option="--cpp_ext" --global-option="--cuda_ext" ./
+```
+
+then set this use_amp to True.
+``` {.python}
+# DEFAULT
+trainer = Trainer(amp_level='O2', use_amp=False)
+```
+
+---
+#### Single-gpu
+Make sure you're on a GPU machine. 
+```python
+# set these flags
+os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
+os.environ["CUDA_VISIBLE_DEVICES"] = "0"
+
+# DEFAULT
+trainer = Trainer(gpus=[0])
+```
+
+---
+#### multi-gpu 
+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"
+
+# DEFAULT
+trainer = Trainer(gpus=[0,1,2,3,4,5,6,7])
+```
+
+---
+#### Multi-node
+COMING SOON.
+
+---
+#### Self-balancing architecture
+Here lightning distributes parts of your module across available GPUs to optimize for speed and memory.   
+
+COMING SOON.

docs/Trainer/Logging.md

@@ -0,0 +1,61 @@
+Lighting offers a few options for logging information about model, gpu usage, etc (via test-tube). It also offers printing options for training monitoring.
+
+
+---
+#### Display metrics in progress bar 
+``` {.python}
+# DEFAULT
+trainer = Trainer(progress_bar=True)
+```
+
+---
+#### Log metric row every k batches 
+Every k batches lightning will make an entry in the metrics log
+``` {.python}
+# DEFAULT (ie: save a .csv log file every 10 batches)
+trainer = Trainer(add_log_row_interval=10)
+```
+
+---
+#### Process position
+When running multiple models on the same machine we want to decide which progress bar to use.
+Lightning will stack progress bars according to this value. 
+``` {.python}
+# DEFAULT
+trainer = Trainer(process_position=0)
+
+# if this is the second model on the node, show the second progress bar below
+trainer = Trainer(process_position=1)
+```
+
+---
+#### Save a snapshot of all hyperparameters 
+Whenever you call .save() on the test-tube experiment it logs all the hyperparameters in current use.
+Give lightning a test-tube Experiment object to automate this for you.
+``` {.python}
+from test-tube import Experiment
+
+exp = Experiment(...)
+Trainer(experiment=exp)
+```
+
+---
+#### Snapshot code for a training run
+Whenever you call .save() on the test-tube experiment it snapshows all code and pushes to a git tag.
+Give lightning a test-tube Experiment object to automate this for you.
+``` {.python}
+from test-tube import Experiment
+
+exp = Experiment(create_git_tag=True)
+Trainer(experiment=exp)
+```
+
+
+---
+#### Write logs file to csv every k batches 
+Every k batches, lightning will write the new logs to disk
+``` {.python}
+# DEFAULT (ie: save a .csv log file every 100 batches)
+trainer = Trainer(log_save_interval=100)
+```
+

docs/Trainer/SLURM Managed Cluster.md

@@ -0,0 +1,104 @@
+Lightning supports model training on a cluster managed by SLURM in the following cases:    
+
+1. Training on single or multi-cpus only.
+2. Training on single or multi-gpus on the same node.
+3. Coming SOON: Training across multiple nodes.
+
+---
+#### Running grid search on a cluster
+To use lightning to run a hyperparameter search (grid-search or random-search) on a cluster do 4 things:   
+
+(1). Define the parameters for the grid search    
+    
+```{.python}
+from test_tube import HyperOptArgumentParser
+
+# subclass of argparse
+parser = HyperOptArgumentParser(strategy='random_search')
+parser.add_argument('--learning_rate', default=0.002, type=float, help='the learning rate')
+
+# let's enable optimizing over the number of layers in the network
+parser.opt_list('--nb_layers', default=2, type=int, tunable=True, options=[2, 4, 8])
+
+hparams = parser.parse_args()    
+```    
+    
+     
+(2). Define the cluster options in the [SlurmCluster object](https://williamfalcon.github.io/test-tube/hpc/SlurmCluster/) (over 5 nodes and 8 gpus)    
+
+```{.python}
+from test_tube.hpc import SlurmCluster
+
+# hyperparameters is a test-tube hyper params object
+# see https://williamfalcon.github.io/test-tube/hyperparameter_optimization/HyperOptArgumentParser/
+hyperparams = args.parse()
+
+# init cluster
+cluster = SlurmCluster(
+    hyperparam_optimizer=hyperparams,
+    log_path='/path/to/log/results/to',
+    python_cmd='python3'
+)
+
+# let the cluster know where to email for a change in job status (ie: complete, fail, etc...)
+cluster.notify_job_status(email='some@email.com', on_done=True, on_fail=True)
+
+# set the job options. In this instance, we'll run 20 different models
+# each with its own set of hyperparameters giving each one 1 GPU (ie: taking up 20 GPUs)
+cluster.per_experiment_nb_gpus = 8
+cluster.per_experiment_nb_nodes = 5
+
+# we'll request 10GB of memory per node
+cluster.memory_mb_per_node = 10000
+
+# set a walltime of 10 minues
+cluster.job_time = '10:00'
+```
+
+(3). Give trainer the cluster_manager in your main function:    
+
+```{.python}
+from pytorch_lightning import Trainer
+
+def train_fx(trial_hparams, cluster_manager, _):
+    # hparams has a specific set of hyperparams
+    
+    my_model = MyLightningModel()
+    
+    # give the trainer the cluster object
+    trainer = Trainer(cluster=cluster_manager)
+    trainer.fit(my_model)
+
+```
+
+(4). Start the grid search     
+```{.python}
+# run the models on the cluster
+cluster.optimize_parallel_cluster_gpu(
+    train_fx, 
+    nb_trials=20, 
+    job_name='my_grid_search_exp_name', 
+    job_display_name='my_exp')
+```
+
+That's it! The SlurmCluster object will automatically checkpoint the lightning model and resubmit if it runs into the walltime!
+
+
+---
+#### Walltime auto-resubmit
+Lightning automatically resubmits jobs when they reach the walltime. You get this behavior for free if you give lightning
+a slurm cluster object.
+
+```{.python}
+def my_main_fx(hparams, slurm_manager, _):
+    trainer = Trainer(cluster=slurm_manager)
+``` 
+
+(See the grid search example above for cluster configuration).
+With this feature lightning will:    
+
+1. automatically checkpoint the model
+2. checkpoint the trainer session
+3. resubmit a continuation job.
+4. load the checkpoint and trainer session in the new model
+

docs/Trainer/Training Loop.md

@@ -0,0 +1,72 @@
+The lightning training loop handles everything except the actual computations of your model. To decide what will happen in your training loop, define the [training_step function](../../Pytorch-lightning/LightningModule/#training_step).
+
+Below are all the things lightning automates for you in the training loop.
+
+---
+#### Accumulated gradients  
+Accumulated gradients runs K small batches of size N before doing a backwards pass. The effect is a large effective batch size of size KxN. 
+
+``` {.python}
+# DEFAULT (ie: no accumulated grads)
+trainer = Trainer(accumulate_grad_batches=1)
+```
+
+---
+#### Anneal Learning rate
+Cut the learning rate by 10 at every epoch listed in this list.
+``` {.python}
+# DEFAULT (don't anneal)
+trainer = Trainer(lr_scheduler_milestones=None)
+
+# cut LR by 10 at 100, 200, and 300 epochs 
+trainer = Trainer(lr_scheduler_milestones=[100, 200, 300])
+```
+
+---
+#### Force training for min or max epochs
+It can be useful to force training for a minimum number of epochs or limit to a max number
+``` {.python}
+# DEFAULT
+trainer = Trainer(min_nb_epochs=1, max_nb_epochs=1000)
+```
+
+---
+#### Force disable early stop 
+Use this to turn off early stopping and run training to the [max_epoch](#force-training-for-min-or-max-epochs)
+``` {.python}
+# DEFAULT
+trainer = Trainer(enable_early_stop=True)
+```
+
+---
+#### Gradient Clipping 
+Use this to turn off early stopping and run training to the [max_epoch](#force-training-for-min-or-max-epochs)
+``` {.python}
+# DEFAULT (ie: don't clip)
+trainer = Trainer(gradient_clip=0)
+```
+
+
+
+---
+#### Inspect gradient norms
+Looking at grad norms can help you figure out where training might be going wrong.
+``` {.python}
+# DEFAULT (-1 doesn't track norms)
+trainer = Trainer(track_grad_norm=-1)
+
+# track the LP norm (P=2 here)
+trainer = Trainer(track_grad_norm=2)
+```
+
+
+---
+#### Set how much of the training set to check
+If you don't want to check 100% of the training set (for debugging or if it's huge), set this flag
+``` {.python}
+# DEFAULT
+trainer = Trainer(train_percent_check=1.0)
+
+# check 10% only
+trainer = Trainer(train_percent_check=0.1)
+```

docs/Trainer/Validation loop.md

@@ -0,0 +1,57 @@
+The lightning validation loop handles everything except the actual computations of your model. To decide what will happen in your validation loop, define the [validation_step function](../../Pytorch-lightning/LightningModule/#validation_step).
+Below are all the things lightning automates for you in the validation loop.
+
+**Note**   
+Lightning will run 5 steps of validation in the beginning of training as a sanity check so you don't have to wait until a full epoch to catch possible validation issues.
+
+
+
+
+---
+#### Check validation every n epochs
+If you have a small dataset you might want to check validation every n epochs
+``` {.python}
+# DEFAULT
+trainer = Trainer(check_val_every_n_epoch=1)
+```
+
+---
+#### Set how much of the validation set to check 
+If you don't want to check 100% of the validation set (for debugging or if it's huge), set this flag
+``` {.python}
+# DEFAULT
+trainer = Trainer(val_percent_check=1.0)
+
+# check 10% only
+trainer = Trainer(val_percent_check=0.1)
+```
+
+---
+#### Set how much of the test set to check 
+If you don't want to check 100% of the test set (for debugging or if it's huge), set this flag
+``` {.python}
+# DEFAULT
+trainer = Trainer(test_percent_check=1.0)
+
+# check 10% only
+trainer = Trainer(test_percent_check=0.1)
+```
+
+---
+####  Set validation check frequency within 1 training epoch
+For large datasets it's often desirable to check validation multiple times within a training loop
+``` {.python}
+# DEFAULT
+trainer = Trainer(val_check_interval=0.95)
+
+# check every .25 of an epoch 
+trainer = Trainer(val_check_interval=0.25)
+```
+
+---
+####  Set the number of validation sanity steps
+Lightning runs a few steps of validation in the beginning of training. This avoids crashing in the validation loop sometime deep into a lengthy training loop.
+``` {.python}
+# DEFAULT
+trainer = Trainer(nb_sanity_val_steps=5)
+```

docs/Trainer/debugging.md

@@ -0,0 +1,48 @@
+These flags are useful to help debug a model.
+
+---
+#### Fast dev run 
+This flag is meant for debugging a full train/val/test loop. It'll activate callbacks, everything but only with 1 training and 1 validation batch.
+Use this to debug a full run of your program quickly
+``` {.python}
+# DEFAULT
+trainer = Trainer(fast_dev_run=False)
+```
+
+---
+#### Inspect gradient norms
+Looking at grad norms can help you figure out where training might be going wrong.
+``` {.python}
+# DEFAULT (-1 doesn't track norms)
+trainer = Trainer(track_grad_norm=-1)
+
+# track the LP norm (P=2 here)
+trainer = Trainer(track_grad_norm=2)
+```
+
+---
+#### Make model overfit on subset of data
+A useful debugging trick is to make your model overfit a tiny fraction of the data.
+``` {.python}
+# DEFAULT don't overfit (ie: normal training)
+trainer = Trainer(overfit_pct=0.0)
+
+# overfit on 1% of data 
+trainer = Trainer(overfit_pct=0.01)
+```
+
+---
+#### Print the parameter count by layer
+By default lightning prints a list of parameters *and submodules* when it starts training.
+
+---
+#### Print which gradients are nan 
+This option prints a list of tensors with nan gradients.
+``` {.python}
+# DEFAULT
+trainer = Trainer(print_nan_grads=False)
+```
+
+---
+#### Log GPU usage
+Lightning automatically logs gpu usage to the test tube logs. It'll only do it at the metric logging interval, so it doesn't slow down training.

docs/Trainer/index.md

@@ -0,0 +1,74 @@
+# Trainer
+[[Github Code](https://github.com/williamFalcon/pytorch-lightning/blob/master/pytorch_lightning/models/trainer.py)]
+
+The lightning trainer abstracts best practices for running a training, val, test routine. It calls parts of your model when it wants to hand over full control and otherwise makes training assumptions which are now standard practice in AI research.
+
+This is the basic use of the trainer:
+
+``` {.python}
+from pytorch_lightning import Trainer
+
+model = LightningTemplate()
+
+trainer = Trainer()
+trainer.fit(model)
+```
+
+But of course the fun is in all the advanced things it can do:
+
+
+**Checkpointing**    
+
+- Model saving
+- Model loading 
+
+**Computing cluster (SLURM)**    
+
+- [Running grid search on a cluster](SLURM%20Managed%20Cluster/#running-grid-search-on-a-cluster)
+- [Walltime auto-resubmit](SLURM%20Managed%20Cluster/#walltime-auto-resubmit)   
+
+**Debugging**  
+
+- [Fast dev run](Debugging/#fast-dev-run)
+- [Inspect gradient norms](Debugging/#inspect-gradient-norms)
+- [Log GPU usage](Debugging/#Log-gpu-usage)
+- [Make model overfit on subset of data](Debugging/#make-model-overfit-on-subset-of-data)
+- [Print the parameter count by layer](Debugging/#print-the-parameter-count-by-layer)
+- [Pring which gradients are nan](Debugging/#print-which-gradients-are-nan)
+
+
+**Distributed training**    
+
+- [16-bit mixed precision](Distributed%20training/#16-bit-mixed-precision)
+- [Multi-GPU](Distributed%20training/#Multi-GPU)
+- [Multi-node](Distributed%20training/#Multi-node)
+- [Single GPU](Distributed%20training/#single-gpu)
+- [Self-balancing architecture](Distributed%20training/#self-balancing-architecture)
+
+
+**Experiment Logging**   
+
+- [Display metrics in progress bar](Logging/#display-metrics-in-progress-bar)
+- Log arbitrary metrics
+- [Log metric row every k batches](Logging/#log-metric-row-every-k-batches)
+- [Process position](Logging/#process-position)
+- [Save a snapshot of all hyperparameters](Logging/#save-a-snapshot-of-all-hyperparameters) 
+- [Snapshot code for a training run](Logging/#snapshot-code-for-a-training-run) 
+- [Write logs file to csv every k batches](Logging/#write-logs-file-to-csv-every-k-batches)
+
+**Training loop**    
+
+- [Accumulate gradients](Training%20Loop/#accumulated-gradients)
+- [Anneal Learning rate](Training%20Loop/#anneal-learning-rate)
+- [Force training for min or max epochs](Training%20Loop/#force-training-for-min-or-max-epochs)
+- [Force disable early stop](Training%20Loop/#force-disable-early-stop)
+- [Use multiple optimizers (like GANs)](../Pytorch-lightning/LightningModule/#configure_optimizers)
+- [Set how much of the training set to check (1-100%)](Training%20Loop/#set-how-much-of-the-training-set-to-check)
+
+**Validation loop**    
+
+- [Check validation every n epochs](Validation%20Loop/#check-validation-every-n-epochs)
+- [Set how much of the validation set to check](Validation%20Loop/#set-how-much-of-the-validation-set-to-check)
+- [Set how much of the test set to check](Validation%20Loop/#set-how-much-of-the-test-set-to-check)
+- [Set validation check frequency within 1 training epoch](Validation%20Loop/#set-validation-check-frequency-within-1-training-epoch)
+- [Set the number of validation sanity steps](Validation%20Loop/#set-the-number-of-validation-sanity-steps)

docs/examples/Examples.md

@@ -0,0 +1,171 @@
+### Template model definition
+In 99% of cases you want to just copy [this template](https://github.com/williamFalcon/pytorch-lightning/blob/master/examples/new_project_templates/lightning_module_template.py) to start a new lightningModule and change the core of what your model is actually trying to do.
+
+```bash
+# get a copy of the module template
+wget https://github.com/williamFalcon/pytorch-lightning/blob/master/examples/new_project_templates/lightning_module_template.py
+```
+
+---    
+### Trainer Example 
+
+** \_\_main__ function**    
+
+Normally, we want to let the \_\_main__ function start the training.
+Inside the main we parse training arguments with whatever hyperparameters we want. Your LightningModule will have a 
+chance to add hyperparameters.   
+
+```{.python}
+from test_tube import HyperOptArgumentParser
+
+if __name__ == '__main__':
+
+    # use default args given by lightning
+    root_dir = os.path.split(os.path.dirname(sys.modules['__main__'].__file__))[0]
+    parent_parser = HyperOptArgumentParser(strategy='random_search', add_help=False)
+    add_default_args(parent_parser, root_dir)
+
+    # allow model to overwrite or extend args
+    parser = ExampleModel.add_model_specific_args(parent_parser)
+    hyperparams = parser.parse_args()
+
+    # train model
+    main(hyperparams)
+```
+**Main Function**      
+
+The main function is your entry into the program. This is where you init your model, checkpoint directory, and launch the training.
+The main function should have 3 arguments:   
+- hparams: a configuration of hyperparameters.    
+- slurm_manager: Slurm cluster manager object (can be None)
+- dict: for you to return any values you want (useful in meta-learning, otherwise set to _)    
+
+```{}
+def main(hparams, cluster, results_dict):
+    """
+    Main training routine specific for this project
+    :param hparams:
+    :return:
+    """
+    # init experiment
+    log_dir = os.path.dirname(os.path.realpath(__file__))
+    exp = Experiment(
+        name='test_tube_exp',
+        debug=True,
+        save_dir=log_dir,
+        version=0,
+        autosave=False,
+        description='test demo'
+    )
+    
+    # set the hparams for the experiment
+    exp.argparse(hparams)
+    exp.save()
+
+    # build model
+    model = MyLightningModule(hparams)
+
+    # callbacks
+    early_stop = EarlyStopping(
+        monitor=hparams.early_stop_metric,
+        patience=hparams.early_stop_patience,
+        verbose=True,
+        mode=hparams.early_stop_mode
+    )
+
+    model_save_path = '{}/{}/{}'.format(hparams.model_save_path, exp.name, exp.version)
+    checkpoint = ModelCheckpoint(
+        filepath=model_save_path,
+        save_function=None,
+        save_best_only=True,
+        verbose=True,
+        monitor=hparams.model_save_monitor_value,
+        mode=hparams.model_save_monitor_mode
+    )
+
+    # configure trainer
+    trainer = Trainer(
+        experiment=exp,
+        cluster=cluster,
+        checkpoint_callback=checkpoint,
+        early_stop_callback=early_stop,
+    )
+
+    # train model
+    trainer.fit(model)
+```
+
+
+
+
+The __main__ function will start training on your **main** function. If you use the HyperParameterOptimizer
+in hyper parameter optimization mode, this main function will get one set of hyperparameters. If you use it as a simple
+argument parser you get the default arguments in the argument parser.
+
+So, calling main(hyperparams) runs the model with the default argparse arguments.       
+```{.python}
+main(hyperparams)
+```
+
+---
+#### CPU hyperparameter search      
+
+```{.python}
+# run a grid search over 20 hyperparameter combinations.
+hyperparams.optimize_parallel_cpu(
+    main_local,
+    nb_trials=20,
+    nb_workers=1
+)
+```
+
+---
+#### Hyperparameter search on a single or multiple GPUs       
+```{.python}
+# run a grid search over 20 hyperparameter combinations.
+hyperparams.optimize_parallel_gpu(
+    main_local,
+    nb_trials=20,
+    nb_workers=1,
+    gpus=[0,1,2,3]
+)
+```
+
+---
+#### Hyperparameter search on a SLURM HPC cluster   
+```{.python}    
+def optimize_on_cluster(hyperparams):
+    # enable cluster training
+    cluster = SlurmCluster(
+        hyperparam_optimizer=hyperparams,
+        log_path=hyperparams.tt_save_path,
+        test_tube_exp_name=hyperparams.tt_name
+    )
+
+    # email for cluster coms
+    cluster.notify_job_status(email='add_email_here', on_done=True, on_fail=True)
+
+    # configure cluster
+    cluster.per_experiment_nb_gpus = hyperparams.per_experiment_nb_gpus
+    cluster.job_time = '48:00:00'
+    cluster.gpu_type = '1080ti'
+    cluster.memory_mb_per_node = 48000
+
+    # any modules for code to run in env
+    cluster.add_command('source activate pytorch_lightning')
+
+    # name of exp
+    job_display_name = hyperparams.tt_name.split('_')[0]
+    job_display_name = job_display_name[0:3]
+
+    # run hopt
+    print('submitting jobs...')
+    cluster.optimize_parallel_cluster_gpu(
+        main,
+        nb_trials=hyperparams.nb_hopt_trials,
+        job_name=job_display_name
+    )
+
+# run cluster hyperparameter search    
+optimize_on_cluster(hyperparams)
+```

docs/index.md

@@ -0,0 +1,78 @@
+# PYTORCH-LIGHTNING DOCUMENTATION
+
+###### 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)
+
+###### Docs shortcuts
+- [LightningModule](LightningModule/RequiredTrainerInterface/)  
+- [Trainer](Trainer/)  
+
+###### Quick start examples 
+- [CPU example](examples/Examples/#cpu-hyperparameter-search)   
+- [Hyperparameter search on single GPU](examples/Examples/#hyperparameter-search-on-a-single-or-multiple-gpus)    
+- [Hyperparameter search on multiple GPUs on same node](examples/Examples/#hyperparameter-search-on-a-single-or-multiple-gpus)  
+- [Hyperparameter search on a SLURM HPC cluster](examples/Examples/#Hyperparameter search on a SLURM HPC cluster)      
+
+
+###### Checkpointing    
+
+- [Model saving](https://williamfalcon.github.io/pytorch-lightning/Trainer/Checkpointing/#model-saving)
+- [Model loading](https://williamfalcon.github.io/pytorch-lightning/LightningModule/methods/#load-from-metrics) 
+
+###### Computing cluster (SLURM)    
+
+- [Running grid search on a cluster](https://williamfalcon.github.io/pytorch-lightning/Trainer/SLURM%20Managed%20Cluster#running-grid-search-on-a-cluster) 
+- [Walltime auto-resubmit](https://williamfalcon.github.io/pytorch-lightning/Trainer/SLURM%20Managed%20Cluster#walltime-auto-resubmit)   
+
+###### Debugging  
+
+- [Fast dev run](https://williamfalcon.github.io/pytorch-lightning/Trainer/debugging/#fast-dev-run)
+- [Inspect gradient norms](https://williamfalcon.github.io/pytorch-lightning/Trainer/debugging/#inspect-gradient-norms)
+- [Log GPU usage](https://williamfalcon.github.io/pytorch-lightning/Trainer/debugging/#Log-gpu-usage)
+- [Make model overfit on subset of data](https://williamfalcon.github.io/pytorch-lightning/Trainer/debugging/#make-model-overfit-on-subset-of-data)
+- [Print the parameter count by layer](https://williamfalcon.github.io/pytorch-lightning/Trainer/debugging/#print-the-parameter-count-by-layer)
+- [Pring which gradients are nan](https://williamfalcon.github.io/pytorch-lightning/Trainer/debugging/#print-which-gradients-are-nan)
+
+
+###### Distributed training    
+
+- [16-bit mixed precision](https://williamfalcon.github.io/pytorch-lightning/Trainer/Distributed%20training/#16-bit-mixed-precision)
+- [Multi-GPU](https://williamfalcon.github.io/pytorch-lightning/Trainer/Distributed%20training/#Multi-GPU)
+- [Multi-node](https://williamfalcon.github.io/pytorch-lightning/Trainer/Distributed%20training/#Multi-node)
+- [Single GPU](https://williamfalcon.github.io/pytorch-lightning/Trainer/Distributed%20training/#single-gpu)
+- [Self-balancing architecture](https://williamfalcon.github.io/pytorch-lightning/Trainer/Distributed%20training/#self-balancing-architecture)
+
+
+###### Experiment Logging   
+
+- [Display metrics in progress bar](https://williamfalcon.github.io/pytorch-lightning/Trainer/Logging/#display-metrics-in-progress-bar)
+- Log arbitrary metrics
+- [Log metric row every k batches](https://williamfalcon.github.io/pytorch-lightning/Trainer/Logging/#log-metric-row-every-k-batches)
+- [Process position](https://williamfalcon.github.io/pytorch-lightning/Trainer/Logging/#process-position)
+- [Save a snapshot of all hyperparameters](https://williamfalcon.github.io/pytorch-lightning/Trainer/Logging/#save-a-snapshot-of-all-hyperparameters) 
+- [Snapshot code for a training run](https://williamfalcon.github.io/pytorch-lightning/Trainer/Logging/#snapshot-code-for-a-training-run) 
+- [Write logs file to csv every k batches](https://williamfalcon.github.io/pytorch-lightning/Trainer/Logging/#write-logs-file-to-csv-every-k-batches)
+
+###### Training loop    
+
+- [Accumulate gradients](https://williamfalcon.github.io/pytorch-lightning/Trainer/Training%20Loop/#accumulated-gradients)
+- [Anneal Learning rate](https://williamfalcon.github.io/pytorch-lightning/Trainer/Training%20Loop/#anneal-learning-rate)
+- [Force training for min or max epochs](https://williamfalcon.github.io/pytorch-lightning/Trainer/Training%20Loop/#force-training-for-min-or-max-epochs)
+- [Force disable early stop](https://williamfalcon.github.io/pytorch-lightning/Trainer/Training%20Loop/#force-disable-early-stop)
+- [Gradient Clipping](https://williamfalcon.github.io/pytorch-lightning/Trainer/Training%20Loop/#gradient-clipping)
+- [Use multiple optimizers (like GANs)](https://williamfalcon.github.io/pytorch-lightning/Pytorch-Lightning/LightningModule/#configure_optimizers)
+- [Set how much of the training set to check (1-100%)](https://williamfalcon.github.io/pytorch-lightning/Trainer/Training%20Loop/#set-how-much-of-the-training-set-to-check)
+
+######Validation loop    
+
+- [Check validation every n epochs](https://williamfalcon.github.io/pytorch-lightning/Trainer/Validation%20loop/#check-validation-every-n-epochs)
+- [Set how much of the validation set to check](https://williamfalcon.github.io/pytorch-lightning/Trainer/Validation%20loop/#set-how-much-of-the-validation-set-to-check)
+- [Set how much of the test set to check](https://williamfalcon.github.io/pytorch-lightning/Trainer/Validation%20loop/#set-how-much-of-the-test-set-to-check)
+- [Set validation check frequency within 1 training epoch](https://williamfalcon.github.io/pytorch-lightning/Trainer/Validation%20loop/#set-validation-check-frequency-within-1-training-epoch)
+- [Set the number of validation sanity steps](https://williamfalcon.github.io/pytorch-lightning/Trainer/Validation%20loop/#set-the-number-of-validation-sanity-steps)
+

examples/new_project_templates/lightning_module_template.py

@@ -1,21 +1,28 @@
+import os
+from collections import OrderedDict
 import torch.nn as nn
-import numpy as np
-from pytorch-lightning.root_module.root_module import RootModule
-from test_tube import HyperOptArgumentParser
 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 pytorch_lightning.root_module.root_module import LightningModule
 
 
-class ExampleModel(RootModule):
+class LightningTemplateModel(LightningModule):
     """
     Sample model to show how to define a template
     """
 
     def __init__(self, hparams):
+        """
+        Pass in parsed HyperOptArgumentParser to the model
+        :param hparams:
+        """
         # init superclass
-        super(ExampleModel, self).__init__(hparams)
+        super(LightningTemplateModel, self).__init__(hparams)
 
         self.batch_size = hparams.batch_size
 
@@ -40,8 +47,14 @@ class ExampleModel(RootModule):
     # 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 = F.tanh(x)
+        x = torch.tanh(x)
         x = self.c_d1_bn(x)
         x = self.c_d1_drop(x)
 
@@ -54,9 +67,9 @@ class ExampleModel(RootModule):
         nll = F.nll_loss(logits, labels)
         return nll
 
-    def training_step(self, data_batch):
+    def training_step(self, data_batch, batch_i):
         """
-        Called inside the training loop
+        Lightning calls this inside the training loop
         :param data_batch:
         :return:
         """
@@ -68,12 +81,15 @@ class ExampleModel(RootModule):
         # calculate loss
         loss_val = self.loss(y, y_hat)
 
-        tqdm_dic = {'tng_loss': loss_val.item()}
+        output = OrderedDict({
-        return loss_val, tqdm_dic
+            'loss': loss_val,
+            'tqdm_metrics': {}
+        })
+        return output
 
-    def validation_step(self, data_batch):
+    def validation_step(self, data_batch, batch_i):
         """
-        Called inside the validation loop
+        Lightning calls this inside the validation loop
         :param data_batch:
         :return:
         """
@@ -87,7 +103,10 @@ class ExampleModel(RootModule):
         labels_hat = torch.argmax(y_hat, dim=1)
         val_acc = torch.sum(y == labels_hat).item() / (len(y) * 1.0)
 
-        output = {'y_hat': y_hat, 'val_loss': loss_val.item(), 'val_acc': val_acc}
+        output = OrderedDict({
+            'val_loss': loss_val,
+            'val_acc': torch.tensor(val_acc),
+        })
         return output
 
     def validation_end(self, outputs):
@@ -97,13 +116,14 @@ class ExampleModel(RootModule):
         :return:
         """
         val_loss_mean = 0
-        accs = []
+        val_acc_mean = 0
         for output in outputs:
             val_loss_mean += output['val_loss']
-            accs.append(output['val_acc'])
+            val_acc_mean += output['val_acc']
 
         val_loss_mean /= len(outputs)
-        tqdm_dic = {'val_loss': val_loss_mean, 'val_acc': np.mean(accs)}
+        val_acc_mean /= len(outputs)
+        tqdm_dic = {'val_loss': val_loss_mean.item(), 'val_acc': val_acc_mean.item()}
         return tqdm_dic
 
     def update_tng_log_metrics(self, logs):
@@ -128,9 +148,8 @@ class ExampleModel(RootModule):
         return whatever optimizers we want here
         :return: list of optimizers
         """
-        optimizer = self.choose_optimizer(self.hparams.optimizer_name, self.parameters(), {'lr': self.hparams.learning_rate}, 'optimizer')
+        optimizer = optim.Adam(self.parameters(), lr=self.hparams.learning_rate)
-        self.optimizers = [optimizer]
+        return [optimizer]
-        return self.optimizers
 
     def __dataloader(self, train):
         # init data generators
@@ -177,7 +196,13 @@ class ExampleModel(RootModule):
         return self._test_dataloader
 
     @staticmethod
-    def add_model_specific_args(parent_parser):
+    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
@@ -186,11 +211,11 @@ class ExampleModel(RootModule):
         # 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)
-        parser.add_argument('--hidden_dim', default=500)
         parser.add_argument('--out_features', default=10)
+        parser.add_argument('--hidden_dim', default=50000) # use 500 for CPU, 50000 for GPU to see speed difference
 
         # data
-        parser.add_argument('--data_root', default='/Users/williamfalcon/Developer/personal/research_lib/research_proj/datasets/mnist', type=str)
+        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, type=float, options=[0.0001, 0.0005, 0.001, 0.005],

examples/new_project_templates/trainer_cpu_template.py

@@ -2,10 +2,10 @@ import os
 import sys
 
 from test_tube import HyperOptArgumentParser, Experiment
-from pytorch-lightning.models.trainer import Trainer
+from pytorch_lightning.models.trainer import Trainer
-from pytorch-lightning.utils.arg_parse import add_default_args
+from pytorch_lightning.utils.arg_parse import add_default_args
-from pytorch-lightning.utils.pt_callbacks import EarlyStopping, ModelCheckpoint
+from pytorch_lightning.callbacks.pt_callbacks import EarlyStopping, ModelCheckpoint
-from demo.example_model import ExampleModel
+from docs.source.examples.example_model import ExampleModel
 
 
 def main(hparams):

examples/new_project_templates/trainer_gpu_cluster_template.py

@@ -5,10 +5,10 @@ from time import sleep
 import torch
 
 from test_tube import HyperOptArgumentParser, Experiment, SlurmCluster
-from pytorch-lightning.models.trainer import Trainer
+from pytorch_lightning.models.trainer import Trainer
-from pytorch-lightning.utils.arg_parse import add_default_args
+from pytorch_lightning.utils.arg_parse import add_default_args
 
-from pytorch-lightning.utils.pt_callbacks import EarlyStopping, ModelCheckpoint
+from pytorch_lightning.callbacks import EarlyStopping, ModelCheckpoint
 
 SEED = 2334
 torch.manual_seed(SEED)
@@ -17,11 +17,11 @@ np.random.seed(SEED)
 # ---------------------
 # DEFINE MODEL HERE
 # ---------------------
-from example_model import ExampleModel
+from examples.new_project_templates.lightning_module_template import LightningTemplateModel
 # ---------------------
 
 AVAILABLE_MODELS = {
-    'model_template': ExampleModel
+    'model_template': LightningTemplateModel
 }
 
 
@@ -42,9 +42,7 @@ def main(hparams, cluster, results_dict):
     :param hparams:
     :return:
     """
-    on_gpu = torch.cuda.is_available()
+    on_gpu = hparams.gpus is not None and torch.cuda.is_available()
-    if hparams.disable_cuda:
-        on_gpu = False
 
     device = 'cuda' if on_gpu else 'cpu'
     hparams.__setattr__('device', device)
@@ -93,12 +91,18 @@ def main(hparams, cluster, results_dict):
         mode=hparams.model_save_monitor_mode
     )
 
+    # gpus are ; separated for inside a node and , within nodes
+    gpu_list = None
+    if hparams.gpus is not None:
+        gpu_list = [int(x) for x in hparams.gpus.split(';')]
+
     # configure trainer
     trainer = Trainer(
         experiment=exp,
         cluster=cluster,
         checkpoint_callback=checkpoint,
         early_stop_callback=early_stop,
+        gpus=gpu_list
     )
 
     # train model
@@ -159,35 +163,40 @@ if __name__ == '__main__':
         model_name = 'model_template'
 
     # use default args
-    root_dir = os.path.split(os.path.dirname(sys.modules['__main__'].__file__))[0]
+    root_dir = os.path.dirname(os.path.realpath(__file__))
     parent_parser = get_default_parser(strategy='random_search', root_dir=root_dir)
 
     # allow model to overwrite or extend args
     TRAINING_MODEL = AVAILABLE_MODELS[model_name]
-    parser = TRAINING_MODEL.add_model_specific_args(parent_parser)
+    parser = TRAINING_MODEL.add_model_specific_args(parent_parser, root_dir)
     hyperparams = parser.parse_args()
 
     # format GPU layout
     os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
-    gpu_ids = hyperparams.gpus.split(';')
 
+    # ---------------------
     # RUN TRAINING
+    # ---------------------
+
+    # cluster and CPU
     if hyperparams.on_cluster:
-        # Gets called when running via HPC cluster
+        # run on HPC cluster
         print('RUNNING ON SLURM CLUSTER')
+        gpu_ids = hyperparams.gpus.split(';')
         os.environ["CUDA_VISIBLE_DEVICES"] = ','.join(gpu_ids)
         optimize_on_cluster(hyperparams)
 
-    elif hyperparams.single_run_gpu:
+    elif hyperparams.gpus is None:
-        # run on 1 gpu
+        # run on cpu
-        print(f'RUNNING 1 TRIAL ON GPU. gpu: {gpu_ids[0]}')
+        print('RUNNING ON CPU')
-        os.environ["CUDA_VISIBLE_DEVICES"] = gpu_ids[0]
         main(hyperparams, None, None)
 
-    elif hyperparams.local or hyperparams.single_run:
+    # single or multiple GPUs on same machine
-        # run 1 trial but on CPU
+    gpu_ids = hyperparams.gpus.split(';')
-        os.environ["CUDA_VISIBLE_DEVICES"] = '0'
+    if hyperparams.interactive:
-        print('RUNNING LOCALLY')
+        # run on 1 gpu
+        print(f'RUNNING INTERACTIVE MODE ON GPUS. gpu ids: {gpu_ids}')
+        os.environ["CUDA_VISIBLE_DEVICES"] = ','.join(gpu_ids)
         main(hyperparams, None, None)
 
     else:

mkdocs.yml

@@ -0,0 +1,10 @@
+site_name: Pytorch lightning Documentation
+theme:
+  name: 'material'
+docs_dir: docs
+repo_url: https://github.com/williamFalcon/pytorch-lightning
+site_dir: 'site'
+site_description: 'Documentation for Pytorch LightningModule, the researcher version of keras.'
+
+dev_addr: '0.0.0.0:8000'
+#google_analytics: ['UA-aasd', 'sitename']

pytorch_lightning/__init__.py

@@ -0,0 +1 @@
+from .models import Trainer

pytorch_lightning/callbacks/__init__.py

@@ -0,0 +1 @@
+from .pt_callbacks import EarlyStopping, ModelCheckpoint

pytorch_lightning/callbacks/pt_callbacks.py

@@ -1,5 +1,6 @@
 import numpy as np
 import os, shutil
+from pytorch_lightning.pt_overrides.override_data_parallel import LightningDataParallel
 
 
 class Callback(object):
@@ -33,6 +34,8 @@ class Callback(object):
         self.params = params
 
     def set_model(self, model):
+        if type(model) is LightningDataParallel:
+            model = model.module
         self.model = model
 
     def on_epoch_begin(self, epoch, logs=None):
@@ -167,12 +170,11 @@ class ModelCheckpoint(Callback):
         period: Interval (number of epochs) between checkpoints.
     """
 
-    def __init__(self, filepath, save_function, monitor='val_loss', verbose=0,
+    def __init__(self, filepath, monitor='val_loss', verbose=0,
                  save_best_only=False, save_weights_only=False,
                  mode='auto', period=1, prefix=''):
         super(ModelCheckpoint, self).__init__()
         self.monitor = monitor
-        self.save_function = save_function
         self.verbose = verbose
         self.filepath = filepath
         self.save_best_only = save_best_only

pytorch_lightning/models/__init__.py

@@ -0,0 +1 @@
+from .trainer import Trainer

pytorch_lightning/models/sample_model_template/model_template.py

@@ -1,6 +1,6 @@
 import torch.nn as nn
 import numpy as np
-from pytorch_lightning.root_module.root_module import RootModule
+from pytorch_lightning.root_module.root_module import LightningModule
 from test_tube import HyperOptArgumentParser
 from torchvision.datasets import MNIST
 import torchvision.transforms as transforms
@@ -8,7 +8,7 @@ import torch
 import torch.nn.functional as F
 
 
-class ExampleModel1(RootModule):
+class ExampleModel1(LightningModule):
     """
     Sample model to show how to define a template
     """

pytorch_lightning/models/trainer.py

@@ -5,6 +5,7 @@ from pytorch_lightning.root_module.memory import get_gpu_memory_map
 import traceback
 from pytorch_lightning.root_module.model_saving import TrainerIO
 from torch.optim.lr_scheduler import MultiStepLR
+from pytorch_lightning.pt_overrides.override_data_parallel import LightningDataParallel
 import pdb
 
 try:
@@ -13,37 +14,53 @@ try:
 except ModuleNotFoundError:
     APEX_AVAILABLE = False
 
+
+def reduce_distributed_output(output, nb_gpus):
+    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,
                  checkpoint_callback, early_stop_callback,
+                 gradient_clip=0,
                  cluster=None,
                  process_position=0,
                  current_gpu_name=0,
-                 on_gpu=False,
+                 gpus=None,
-                 enable_tqdm=True,
+                 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,
-                 enable_early_stop=True, max_nb_epochs=5, min_nb_epochs=1,
+                 enable_early_stop=True, 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=1, add_log_row_interval=1,
+                 log_save_interval=100, add_log_row_interval=10,
                  lr_scheduler_milestones=None,
                  use_amp=False,
-                 check_grad_nans=False,
+                 print_nan_grads=False,
                  amp_level='O2',
                  nb_sanity_val_steps=5):
 
         # Transfer params
+        self.gradient_clip = gradient_clip
         self.check_val_every_n_epoch = check_val_every_n_epoch
         self.enable_early_stop = enable_early_stop
         self.track_grad_norm = track_grad_norm
         self.fast_dev_run = fast_dev_run
-        self.on_gpu = on_gpu
+        self.on_gpu = gpus is not None and torch.cuda.is_available()
-        self.enable_tqdm = enable_tqdm
+        self.progress_bar = progress_bar
         self.experiment = experiment
         self.exp_save_path = experiment.get_data_path(experiment.name, experiment.version)
         self.cluster = cluster
@@ -61,7 +78,9 @@ class Trainer(TrainerIO):
         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.check_grad_nans = check_grad_nans
+        self.print_nan_grads = print_nan_grads
+        self.data_parallel_device_ids = gpus
+        self.data_parallel = gpus is not None and len(gpus) > 0
 
         # training state
         self.optimizers = None
@@ -115,9 +134,10 @@ class Trainer(TrainerIO):
             'batch_nb':'{}'.format(self.batch_nb),
         }
         tqdm_dic.update(self.tqdm_metrics)
+
         return tqdm_dic
 
-    def __layout_bookeeping(self):
+    def __layout_bookeeping(self, model):
         # training bookeeping
         self.total_batch_nb = 0
         self.running_loss = []
@@ -126,17 +146,17 @@ class Trainer(TrainerIO):
         self.tqdm_metrics = {}
 
         # determine number of training batches
-        self.nb_tng_batches = self.model.nb_batches(self.tng_dataloader)
+        self.nb_tng_batches = model.nb_batches(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 = self.model.nb_batches(self.val_dataloader)
+        self.nb_val_batches = model.nb_batches(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 = self.model.nb_batches(self.test_dataloader)
+        self.nb_test_batches = model.nb_batches(self.test_dataloader)
         self.nb_test_batches = int(self.nb_test_batches * self.test_percent_check)
 
         # determine when to check validation
@@ -144,6 +164,9 @@ class Trainer(TrainerIO):
 
     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):
@@ -159,6 +182,7 @@ class Trainer(TrainerIO):
         # enable eval mode
         model.zero_grad()
         model.eval()
+        model.from_lightning = True
 
         # disable gradients to save memory
         torch.set_grad_enabled(False)
@@ -179,14 +203,22 @@ class Trainer(TrainerIO):
             # -----------------
             # RUN VALIDATION STEP
             # -----------------
+            if self.data_parallel:
+                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.enable_tqdm and self.prog_bar is not None:
+            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
@@ -194,6 +226,7 @@ class Trainer(TrainerIO):
 
         # enable gradients to save memory
         torch.set_grad_enabled(True)
+
         return val_results
 
     def __get_dataloaders(self, model):
@@ -210,14 +243,14 @@ class Trainer(TrainerIO):
     # MODEL TRAINING
     # -----------------------------
     def fit(self, model):
-        self.model = model
+
         model.trainer = self
 
         # transfer data loaders from model
         self.__get_dataloaders(model)
 
         # init training constants
-        self.__layout_bookeeping()
+        self.__layout_bookeeping(model)
 
         # CHOOSE OPTIMIZER
         # filter out the weights that were done on gpu so we can load on good old cpus
@@ -225,8 +258,8 @@ class Trainer(TrainerIO):
 
         if self.use_amp:
             # An example
-            self.model, optimizer = amp.initialize(
+            model, optimizer = amp.initialize(
-                self.model, self.optimizers[0], opt_level=self.amp_level,
+                model, self.optimizers[0], opt_level=self.amp_level,
             )
             self.optimizers[0] = optimizer
             model.trainer = self
@@ -242,7 +275,7 @@ class Trainer(TrainerIO):
 
         # put on gpu if needed
         if self.on_gpu:
-            model = model.cuda()
+            model = LightningDataParallel(model, device_ids=self.data_parallel_device_ids)
 
         # 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)
@@ -257,6 +290,7 @@ class Trainer(TrainerIO):
         # ---------------------------
         # CORE TRAINING LOOP
         # ---------------------------
+        self.model = model
         self.__train()
 
     def __train(self):
@@ -266,24 +300,28 @@ class Trainer(TrainerIO):
             for lr_scheduler in self.lr_schedulers:
                 lr_scheduler.step()
 
-            self.model.current_epoch = epoch_nb
+            model = self.model.module if self.data_parallel else self.model
+            model.current_epoch = epoch_nb
 
             # hook
             if self.__is_function_implemented('on_epoch_start'):
-                self.model.on_epoch_start()
+                model = self.model.module if self.data_parallel else self.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.enable_tqdm:
+            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
-                self.model.global_step = self.global_step
+
+                model = self.model.module if self.data_parallel else self.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
@@ -313,6 +351,9 @@ class Trainer(TrainerIO):
                     # count items in memory
                     # nb_params, nb_tensors = count_mem_items()
 
+                    if self.data_parallel:
+                        metrics = self.model.module.update_tng_log_metrics(self.__tng_tqdm_dic)
+                    else:
                         metrics = self.model.update_tng_log_metrics(self.__tng_tqdm_dic)
 
                     # add gpu memory
@@ -322,7 +363,9 @@ class Trainer(TrainerIO):
 
                     # add norms
                     if self.track_grad_norm > 0:
-                        grad_norm_dic = self.model.grad_norm(self.track_grad_norm)
+                        model = self.model.module if self.data_parallel else self.model
+                        grad_norm_dic = model.grad_norm(self.track_grad_norm)
+
                         metrics.update(grad_norm_dic)
 
                     # log metrics
@@ -331,7 +374,8 @@ class Trainer(TrainerIO):
 
                 # hook
                 if self.__is_function_implemented('on_batch_end'):
-                    self.model.on_batch_end()
+                    model = self.model.module if self.data_parallel else self.model
+                    model.on_batch_end()
 
                 # end epoch early
                 if early_stop_epoch:
@@ -339,7 +383,8 @@ class Trainer(TrainerIO):
 
             # hook
             if self.__is_function_implemented('on_epoch_end'):
-                self.model.on_epoch_end()
+                model = self.model.module if self.data_parallel else self.model
+                model.on_epoch_end()
 
             # early stopping
             if self.enable_early_stop:
@@ -358,16 +403,26 @@ class Trainer(TrainerIO):
 
         # hook
         if self.__is_function_implemented('on_batch_start'):
-            response = self.model.on_batch_start(data_batch)
+            model = self.model.module if self.data_parallel else self.model
+            response = model.on_batch_start(data_batch)
+
             if response == -1:
                 return -1
 
-        if self.enable_tqdm:
+        if self.progress_bar:
             self.prog_bar.update(1)
 
         # forward pass
         # return a scalar value and a dic with tqdm metrics
-        loss, model_specific_tqdm_metrics_dic = self.model.training_step(data_batch, batch_nb)
+        if self.data_parallel:
+            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)
+
+        model_specific_tqdm_metrics_dic = output['tqdm_metrics']
+        loss = output['loss']
+
         self.__add_tqdm_metrics(model_specific_tqdm_metrics_dic)
 
         # backward pass
@@ -378,8 +433,9 @@ class Trainer(TrainerIO):
         else:
             loss.backward()
 
-        if self.check_grad_nans:
+        if self.print_nan_grads:
-            for param in self.model.parameters():
+            model = self.model.module if self.data_parallel else self.model
+            for param in model.parameters():
                 print(param.grad.float().sum())
 
         self.batch_loss_value += loss.item()
@@ -387,6 +443,11 @@ class Trainer(TrainerIO):
         # gradient update with accumulated gradients
         if (self.batch_nb + 1) % self.accumulate_grad_batches == 0:
 
+            # clip gradients
+            if self.gradient_clip > 0:
+                model = self.model.module if self.data_parallel else self.model
+                torch.nn.utils.clip_grad_norm(model.parameters(), self.gradient_clip)
+
             # update gradients across all optimizers
             for optimizer in self.optimizers:
                 optimizer.step()
@@ -403,7 +464,7 @@ class Trainer(TrainerIO):
             self.avg_loss = np.mean(self.running_loss[-100:])
 
             # update progbar
-            if self.enable_tqdm:
+            if self.progress_bar:
                 # add model specific metrics
                 tqdm_metrics = self.__tng_tqdm_dic
                 self.prog_bar.set_postfix(**tqdm_metrics)
@@ -445,7 +506,7 @@ class Trainer(TrainerIO):
             print(e)
             print(traceback.print_exc())
 
-        if self.enable_tqdm:
+        if self.progress_bar:
             # add model specific metrics
             tqdm_metrics = self.__tng_tqdm_dic
             self.prog_bar.set_postfix(**tqdm_metrics)

pytorch_lightning/pt_overrides/override_data_parallel.py

@@ -0,0 +1,105 @@
+from torch.nn import DataParallel
+
+import threading
+import torch
+from torch.cuda._utils import _get_device_index
+import pdb
+
+
+def get_a_var(obj):
+    if isinstance(obj, torch.Tensor):
+        return obj
+
+    if isinstance(obj, list) or isinstance(obj, tuple):
+        for result in map(get_a_var, obj):
+            if isinstance(result, torch.Tensor):
+                return result
+    if isinstance(obj, dict):
+        for result in map(get_a_var, obj.items()):
+            if isinstance(result, torch.Tensor):
+                return result
+    return None
+
+
+class LightningDataParallel(DataParallel):
+    """
+    Override the forward call in lightning so it goes to training and validation step respectively
+    """
+
+    def parallel_apply(self, replicas, inputs, kwargs):
+        return parallel_apply(replicas, inputs, kwargs, self.device_ids[:len(replicas)])
+
+
+def parallel_apply(modules, inputs, kwargs_tup=None, devices=None):
+    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`.
+
+    Args:
+        modules (Module): modules to be parallelized
+        inputs (tensor): inputs to the modules
+        devices (list of int or torch.device): CUDA devices
+
+    :attr:`modules`, :attr:`inputs`, :attr:`kwargs_tup` (if given), and
+    :attr:`devices` (if given) should all have same length. Moreover, each
+    element of :attr:`inputs` can either be a single object as the only argument
+    to a module, or a collection of positional arguments.
+    """
+    assert len(modules) == len(inputs)
+    if kwargs_tup is not None:
+        assert len(modules) == len(kwargs_tup)
+    else:
+        kwargs_tup = ({},) * len(modules)
+    if devices is not None:
+        assert len(modules) == len(devices)
+    else:
+        devices = [None] * len(modules)
+    devices = list(map(lambda x: _get_device_index(x, True), devices))
+    lock = threading.Lock()
+    results = {}
+    grad_enabled = torch.is_grad_enabled()
+
+    def _worker(i, module, input, kwargs, device=None):
+        torch.set_grad_enabled(grad_enabled)
+        if device is None:
+            device = get_a_var(input).get_device()
+        try:
+            with torch.cuda.device(device):
+                # this also avoids accidental slicing of `input` if it is a Tensor
+                if not isinstance(input, (list, tuple)):
+                    input = (input,)
+
+                # ---------------
+                # CHANGE
+                if module.training:
+                    output = module.training_step(*input, **kwargs)
+                else:
+                    output = module.validation_step(*input, **kwargs)
+                # ---------------
+
+            with lock:
+                results[i] = output
+        except Exception as e:
+            with lock:
+                results[i] = e
+
+    if len(modules) > 1:
+        threads = [threading.Thread(target=_worker,
+                                    args=(i, module, input, kwargs, device))
+                   for i, (module, input, kwargs, device) in
+                   enumerate(zip(modules, inputs, kwargs_tup, devices))]
+
+        for thread in threads:
+            thread.start()
+        for thread in threads:
+            thread.join()
+    else:
+        _worker(0, modules[0], inputs[0], kwargs_tup[0], devices[0])
+
+    outputs = []
+    for i in range(len(inputs)):
+        output = results[i]
+        if isinstance(output, Exception):
+            raise output
+        outputs.append(output)
+    return outputs

pytorch_lightning/root_module/model_saving.py

@@ -1,7 +1,8 @@
 import torch
 import os
 import re
-
+import pdb
+from pytorch_lightning.pt_overrides.override_data_parallel import LightningDataParallel
 
 class ModelIO(object):
 
@@ -48,7 +49,8 @@ class TrainerIO(object):
         checkpoint['optimizer_states'] = optimizer_states
 
         # request what to save from the model
-        checkpoint_dict = self.model.get_save_dict()
+        model = self.model.module if type(self.model) is LightningDataParallel else self.model
+        checkpoint_dict = model.get_save_dict()
 
         # merge trainer and model saving items
         checkpoint.update(checkpoint_dict)
@@ -99,6 +101,9 @@ class TrainerIO(object):
     # PRIVATE OPS
     # ----------------------------------
     def hpc_save(self, folderpath, experiment):
+        # make sure the checkpoint folder exists
+        os.makedirs(folderpath, exist_ok=True)
+
         # save exp to make sure we get all the metrics
         experiment.save()
 
@@ -126,10 +131,15 @@ class TrainerIO(object):
         self.restore_training_state(checkpoint)
 
         # load model state
-        self.model.load_model_specific(checkpoint)
+        model = self.model.module if type(self.model) is LightningDataParallel else self.model
+        model.load_model_specific(checkpoint)
 
     def max_ckpt_in_folder(self, path):
         files = os.listdir(path)
+        files = [x for x in files if 'ckpt_' in x]
+        if len(files) == 0:
+            return 0
+
         ckpt_vs = []
         for name in files:
             name = name.split('ckpt_')[-1]

pytorch_lightning/root_module/root_module.py

@@ -9,10 +9,11 @@ from pytorch_lightning.root_module.optimization import OptimizerConfig
 from pytorch_lightning.root_module.hooks import ModelHooks
 
 
-class RootModule(GradInformation, ModelIO, OptimizerConfig, ModelHooks):
+
+class LightningModule(GradInformation, ModelIO, OptimizerConfig, ModelHooks):
 
     def __init__(self, hparams):
-        super(RootModule, self).__init__()
+        super(LightningModule, self).__init__()
         self.hparams = hparams
 
         self.dtype = torch.FloatTensor
@@ -23,8 +24,8 @@ class RootModule(GradInformation, ModelIO, OptimizerConfig, ModelHooks):
         self.fast_dev_run = hparams.fast_dev_run
         self.overfit = hparams.overfit
         self.gradient_clip = hparams.gradient_clip
-        self.num = 2
         self.trainer = None
+        self.from_lightning = True
 
         # track if gpu was requested for checkpointing
         self.on_gpu = False

pytorch_lightning/trainer_main.py

@@ -8,7 +8,7 @@ from pytorch_lightning.models.trainer import Trainer
 from pytorch_lightning.utils.arg_parse import add_default_args
 from time import sleep
 
-from pytorch_lightning.utils.pt_callbacks import EarlyStopping, ModelCheckpoint
+from pytorch_lightning.callbacks.pt_callbacks import EarlyStopping, ModelCheckpoint
 SEED = 2334
 torch.manual_seed(SEED)
 np.random.seed(SEED)
@@ -97,7 +97,7 @@ def main(hparams, cluster, results_dict):
         experiment=exp,
         on_gpu=on_gpu,
         cluster=cluster,
-        enable_tqdm=hparams.enable_tqdm,
+        progress_bar=hparams.enable_tqdm,
         overfit_pct=hparams.overfit,
         track_grad_norm=hparams.track_grad_norm,
         fast_dev_run=hparams.fast_dev_run,

pytorch_lightning/utils/arg_parse.py

@@ -1,3 +1,5 @@
+import pdb
+
 def add_default_args(parser, root_dir, rand_seed=None, possible_model_names=None):
 
     # tng, test, val check intervals
@@ -45,10 +47,8 @@ def add_default_args(parser, root_dir, rand_seed=None, possible_model_names=None
     parser.add_argument('--log_stdout', dest='log_stdout', action='store_true')
 
     # GPU
-    parser.add_argument('--per_experiment_nb_gpus', default=1, type=int)
+    parser.add_argument('--gpus', default=None, type=str)
-    parser.add_argument('--gpus', default='0', type=str)
     parser.add_argument('--single_run_gpu', dest='single_run_gpu', action='store_true')
-    parser.add_argument('--disable_cuda', dest='disable_cuda', action='store_true')
     parser.add_argument('--default_tensor_type', default='torch.cuda.FloatTensor', type=str)
     parser.add_argument('--use_amp', dest='use_amp', action='store_true')
     parser.add_argument('--check_grad_nans', dest='check_grad_nans', action='store_true')
@@ -68,9 +68,9 @@ def add_default_args(parser, root_dir, rand_seed=None, possible_model_names=None
     if rand_seed is not None:
         parser.add_argument('--random_seed', default=rand_seed, type=int)
 
-    parser.add_argument('--live', dest='live', action='store_true', help='runs on gpu without cluster')
+    parser.add_argument('--interactive', dest='interactive', action='store_true', help='runs on gpu without cluster')
-    parser.add_argument('--enable_debug', dest='debug', action='store_true', help='enables/disables test tube')
+    parser.add_argument('--debug', dest='debug', action='store_true', help='enables/disables test tube')
-    parser.add_argument('--enable_local', dest='local', action='store_true', help='enables local tng')
+    parser.add_argument('--local', dest='local', action='store_true', help='enables local tng')
 
     # optimizer
     parser.add_argument('--lr_scheduler_milestones', default=None, type=str)

setup.py

@@ -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.1.dev1832',
+    version='0.112',
     description="The Keras for ML researchers using PyTorch",
     author="William Falcon",
     author_email="waf2107@columbia.edu",
@@ -17,7 +17,7 @@ setup(
     keywords=["deep learning", "pytorch", "AI"],
     python_requires=">=3.5",
     install_requires=[
-        "torch",
+        "torch>=1.0.0",
         "tqdm",
         "test-tube",
     ],