release v0.2.2

.gitignore

@@ -8,6 +8,7 @@ datasets/
 model_weights/
 app/models/
 pip-wheel-metadata/
+test_tube_exp/
 
 # Byte-compiled / optimized / DLL files
 __pycache__/
@@ -116,3 +117,6 @@ ENV/
 
 # mypy
 .mypy_cache/
+
+# data
+mnist/

README.md

@@ -20,221 +20,214 @@ pip install pytorch-lightning
 ```
 
 ## Docs   
-In progress. Documenting now!  
-
-## 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).
+**[View the docs here](https://williamfalcon.github.io/pytorch-lightning/)**
 
 ## 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).   
-2. [Define a model](https://github.com/williamFalcon/pytorch-lightning/blob/master/docs/source/examples/example_model.py).     
+1. [Define a Trainer](https://github.com/williamFalcon/pytorch-lightning/blob/master/examples/new_project_templates/trainer_cpu_template.py).   
+2. [Define a LightningModel](https://github.com/williamFalcon/pytorch-lightning/blob/master/examples/new_project_templates/lightning_module_template.py).     
 
-#### Basic trainer example  
-See [this demo](https://github.com/williamFalcon/pytorch-lightning/blob/master/docs/source/examples/fully_featured_trainer.py) for a more robust trainer example.   
+## What does lightning control for me?
+Everything!    
+Except for these 6 core functions which you define:    
+
+```{.python}
+# what to do in the training loop
+def training_step(self, data_batch, batch_nb):
+
+# what to do in the validation loop
+def validation_step(self, data_batch, batch_nb):
+
+# how to aggregate validation_step outputs
+def validation_end(self, outputs):
+
+# and your dataloaders
+def tng_dataloader():
+def val_dataloader():
+def test_dataloader():
+```
+
+**Could be as complex as seq-2-seq + attention**    
 
 ```python
-import os
-import sys
+# define what happens for training here
+def training_step(self, data_batch, batch_nb):
+    x, y = data_batch
+    
+    # define your own forward and loss calculation
+    hidden_states = self.encoder(x)
+     
+    # even as complex as a seq-2seq + attn model
+    # (this is just a toy, non-working example to illustrate)
+    start_token = '<SOS>'
+    last_hidden = torch.zeros(...)
+    loss = 0
+    for step in range(max_seq_len):
+        attn_context = self.attention_nn(hidden_states, start_token)
+        pred = self.decoder(start_token, attn_context, last_hidden) 
+        last_hidden = pred
+        pred = self.predict_nn(pred)
+        loss += self.loss(last_hidden, y[step])
+        
+    #toy example as well
+    loss = loss / max_seq_len
+    return {'loss': loss} 
+```
 
-from test_tube import HyperOptArgumentParser, Experiment
-from pytorch_lightning.models.trainer import Trainer
-from pytorch_lightning.utils.arg_parse import add_default_args
-from pytorch_lightning.utils.pt_callbacks import EarlyStopping, ModelCheckpoint
-from demo.example_model import ExampleModel
+**Or as basic as CNN image classification**      
 
+```python
+# define what happens for validation here
+def validation_step(self, data_batch, batch_nb):    
+    x, y = data_batch
+    
+    # or as basic as a CNN classification
+    out = self.forward(x)
+    loss = my_loss(out, y)
+    return {'loss': loss} 
+```
 
-def main(hparams):
+**And you also decide how to collate the output of all validation steps**    
+
+```python
+def validation_end(self, outputs):
     """
-    Main training routine specific for this project
-    :param hparams:
+    Called at the end of validation to aggregate outputs
+    :param outputs: list of individual outputs of each validation step
     :return:
     """
-    # init experiment
-    exp = Experiment(
-        name=hparams.tt_name,
-        debug=hparams.debug,
-        save_dir=hparams.tt_save_path,
-        version=hparams.hpc_exp_number,
-        autosave=False,
-        description=hparams.tt_description
-    )
+    val_loss_mean = 0
+    val_acc_mean = 0
+    for output in outputs:
+        val_loss_mean += output['val_loss']
+        val_acc_mean += output['val_acc']
 
-    exp.argparse(hparams)
-    exp.save()
+    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
+```
+   
+## TensorboardX    
+Lightning is fully integrated with tensorboardX.   
 
-    model_save_path = '{}/{}/{}'.format(hparams.model_save_path, exp.name, exp.version)
+<p align="center">
+  <a href="https://williamfalcon.github.io/pytorch-lightning/">
+    <img alt="" src="https://github.com/williamFalcon/pytorch-lightning/blob/master/docs/source/_static/tf_loss.png" width="900px">
+  </a>
+</p>
 
-    # build model
-    model = ExampleModel(hparams)
+Lightning also adds a text column with all the hyperparameters for this experiment.      
 
-    # 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')
+<p align="center">
+  <a href="https://williamfalcon.github.io/pytorch-lightning/">
+        <img alt="" src="https://github.com/williamFalcon/pytorch-lightning/blob/master/docs/source/_static/tf_tags.png" width="900px">
+  </a>
+</p>
 
-    # configure trainer
-    trainer = Trainer(experiment=exp, checkpoint_callback=checkpoint, early_stop_callback=early_stop)
+Simply note the path you set for the Experiment    
+``` {.python}   
+from test_tube import Experiment
+from pytorch-lightning import  Trainer
 
-    # train model
-    trainer.fit(model)
+exp = Experiment(save_dir='/some/path')
+trainer = Trainer(experiment=exp)
+...
+```   
+
+And run tensorboard from that dir   
+```bash
+tensorboard --logdir /some/path     
+```    
+
+## Lightning automatically automates all of the following ([each is also configurable](https://williamfalcon.github.io/pytorch-lightning/Trainer/)):
+
+###### 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)
 
 
-if __name__ == '__main__':
+###### Distributed training    
 
-    # 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)
+- [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)
 
-    # allow model to overwrite or extend args
-    parser = ExampleModel.add_model_specific_args(parent_parser)
-    hyperparams = parser.parse_args()
 
-    # train model
-    main(hyperparams)
+###### 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
 ```
 
-#### Basic model example  
-Here we only show the method signatures. It's up to you to define the content.   
+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"
 
-```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())
-
-        # return a dict
-        return {'total_acc': np.mean(total_accs)}
-
-     # ---------------
-     # 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
-
-    def load_model_specific(self, checkpoint):
-        # lightning loads for you. Here's your chance to say what you want to load
-        self.load_state_dict(checkpoint['state_dict'])
-
-    # ---------------
-    # 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
-    def tng_dataloader(self):
-        return pytorch_dataloader('train')
-
-    @property
-    def val_dataloader(self):
-        return pytorch_dataloader('val')
-
-    @property
-    def test_dataloader(self):
-        return pytorch_dataloader('test')
-
-    # ---------------
-    # MODIFY YOUR COMMAND LINE ARGS
-    @staticmethod
-    def add_model_specific_args(parent_parser):    
-        parser = HyperOptArgumentParser(strategy=parent_parser.strategy, parents=[parent_parser])
-        parser.add_argument('--out_features', default=20)
-        return parser
+# 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,40 @@
+A LightningModule has the following properties which you can access at any time
+
+--- 
+#### current_epoch
+The current epoch   
+
+---
+#### dtype    
+Current dtype    
+
+---
+#### experiment    
+An instance of test-tube Experiment which you can use to log anything for tensorboarX.   
+```{.python}
+self.experiment.add_embedding(...)   
+self.experiment.log({'val_loss': 0.9})   
+self.experiment.add_scalars(...)   
+```
+
+--- 
+#### 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.
+```{.python}   
+self.trainer.optimizers
+self.trainer.current_epoch
+...   
+```   
+

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,76 @@
+###### 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)
+

docs/source/examples/fully_featured_trainer.py

@@ -1,201 +0,0 @@
-import os
-import sys
-import numpy as np
-from time import sleep
-import torch
-
-from test_tube import HyperOptArgumentParser, Experiment, SlurmCluster
-from pytorch-lightning.models.trainer import Trainer
-from pytorch-lightning.utils.arg_parse import add_default_args
-
-from pytorch-lightning.utils.pt_callbacks import EarlyStopping, ModelCheckpoint
-
-SEED = 2334
-torch.manual_seed(SEED)
-np.random.seed(SEED)
-
-# ---------------------
-# DEFINE MODEL HERE
-# ---------------------
-from example_model import ExampleModel
-# ---------------------
-
-AVAILABLE_MODELS = {
-    'model_template': ExampleModel
-}
-
-
-"""
-Allows training by using command line arguments
-Run by:
-# TYPE YOUR RUN COMMAND HERE
-"""
-
-
-def main_local(hparams):
-    main(hparams, None, None)
-
-
-def main(hparams, cluster, results_dict):
-    """
-    Main training routine specific for this project
-    :param hparams:
-    :return:
-    """
-    on_gpu = torch.cuda.is_available()
-    if hparams.disable_cuda:
-        on_gpu = False
-
-    device = 'cuda' if on_gpu else 'cpu'
-    hparams.__setattr__('device', device)
-    hparams.__setattr__('on_gpu', on_gpu)
-    hparams.__setattr__('nb_gpus', torch.cuda.device_count())
-    hparams.__setattr__('inference_mode', hparams.model_load_weights_path is not None)
-
-    # delay each training start to not overwrite logs
-    process_position, current_gpu = TRAINING_MODEL.get_process_position(hparams.gpus)
-    sleep(process_position + 1)
-
-    # init experiment
-    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'
-    )
-
-    exp.argparse(hparams)
-    exp.save()
-
-    # build model
-    print('loading model...')
-    model = TRAINING_MODEL(hparams)
-    print('model built')
-
-    # 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)
-
-
-def get_default_parser(strategy, root_dir):
-
-    possible_model_names = list(AVAILABLE_MODELS.keys())
-    parser = HyperOptArgumentParser(strategy=strategy, add_help=False)
-    add_default_args(parser, root_dir, possible_model_names=possible_model_names, rand_seed=SEED)
-    return parser
-
-
-def get_model_name(args):
-    for i, arg in enumerate(args):
-        if 'model_name' in arg:
-            return args[i+1]
-
-
-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
-    )
-
-
-if __name__ == '__main__':
-
-    model_name = get_model_name(sys.argv)
-    if model_name is None:
-        model_name = 'model_template'
-
-    # use default args
-    root_dir = os.path.split(os.path.dirname(sys.modules['__main__'].__file__))[0]
-    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)
-    hyperparams = parser.parse_args()
-
-    # format GPU layout
-    os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
-    gpu_ids = hyperparams.gpus.split(';')
-
-    # RUN TRAINING
-    if hyperparams.on_cluster:
-        # Gets called when running via HPC cluster
-        print('RUNNING ON SLURM CLUSTER')
-        os.environ["CUDA_VISIBLE_DEVICES"] = ','.join(gpu_ids)
-        optimize_on_cluster(hyperparams)
-
-    elif hyperparams.single_run_gpu:
-        # run on 1 gpu
-        print(f'RUNNING 1 TRIAL ON GPU. gpu: {gpu_ids[0]}')
-        os.environ["CUDA_VISIBLE_DEVICES"] = gpu_ids[0]
-        main(hyperparams, None, None)
-
-    elif hyperparams.local or hyperparams.single_run:
-        # run 1 trial but on CPU
-        os.environ["CUDA_VISIBLE_DEVICES"] = '0'
-        print('RUNNING LOCALLY')
-        main(hyperparams, None, None)
-
-    else:
-        # multiple GPUs on same machine
-        print(f'RUNNING MULTI GPU. GPU ids: {gpu_ids}')
-        hyperparams.optimize_parallel_gpu(
-            main_local,
-            gpu_ids=gpu_ids,
-            nb_trials=hyperparams.nb_hopt_trials,
-            nb_workers=len(gpu_ids)
-        )

examples/new_project_templates/__init__.py

@@ -0,0 +1 @@
+from .lightning_module_template import LightningTemplateModel

examples/new_project_templates/lightning_module_template.py

@@ -1,21 +1,30 @@
+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 torch.utils.data import DataLoader
+from torch.utils.data.distributed import DistributedSampler
+
+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 +49,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 +69,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 +83,15 @@ class ExampleModel(RootModule):
         # calculate loss
         loss_val = self.loss(y, y_hat)
 
-        tqdm_dic = {'tng_loss': loss_val.item()}
-        return loss_val, tqdm_dic
+        output = OrderedDict({
+            '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 +105,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 +118,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,20 +150,31 @@ 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')
-        self.optimizers = [optimizer]
-        return self.optimizers
+        optimizer = optim.Adam(self.parameters(), lr=self.hparams.learning_rate)
+        return [optimizer]
 
     def __dataloader(self, train):
         # init data generators
         transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,), (1.0,))])
-
         dataset = MNIST(root=self.hparams.data_root, train=train, transform=transform, download=True)
 
-        loader = torch.utils.data.DataLoader(
+        # when using multi-node we need to add the datasampler
+        train_sampler = None
+        batch_size = self.hparams.batch_size
+
+        try:
+            if self.on_gpu:
+                train_sampler = DistributedSampler(dataset, rank=self.trainer.proc_rank)
+                batch_size = batch_size // self.trainer.world_size  # scale batch size
+        except Exception as e:
+            pass
+
+        should_shuffle = train_sampler is None
+        loader = DataLoader(
             dataset=dataset,
-            batch_size=self.hparams.batch_size,
-            shuffle=True
+            batch_size=batch_size,
+            shuffle=should_shuffle,
+            sampler=train_sampler
         )
 
         return loader
@@ -177,7 +210,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
@@ -185,16 +224,19 @@ 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('--in_features', default=28*28, type=int)
+        parser.add_argument('--out_features', default=10, type=int)
+        parser.add_argument('--hidden_dim', default=50000, type=int) # use 500 for CPU, 50000 for GPU to see speed difference
 
         # data
-        parser.add_argument('--data_root', default='/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],
+        parser.opt_list('--learning_rate', default=0.001*8, type=float, options=[0.0001, 0.0005, 0.001, 0.005],
                         tunable=False)
-        parser.opt_list('--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)
+
+        # if using 2 nodes with 4 gpus each the batch size here (256) will be 256 / (2*8) = 16 per gpu
+        parser.opt_list('--batch_size', default=256*8, type=int, options=[32, 64, 128, 256], tunable=False,
+                        help='batch size will be divided over all the gpus being used across all nodes')
         return parser

examples/new_project_templates/multi_node_cluster_template.py

@@ -0,0 +1,172 @@
+import os
+import sys
+import numpy as np
+from time import sleep
+import torch
+
+from test_tube import HyperOptArgumentParser, Experiment, SlurmCluster
+from pytorch_lightning.models.trainer import Trainer
+from pytorch_lightning.utils.arg_parse import add_default_args
+
+from pytorch_lightning.callbacks import EarlyStopping, ModelCheckpoint
+
+SEED = 2334
+torch.manual_seed(SEED)
+np.random.seed(SEED)
+
+# ---------------------
+# DEFINE MODEL HERE
+# ---------------------
+from lightning_module_template import LightningTemplateModel
+# ---------------------
+
+"""
+Allows training by using command line arguments
+Run by:
+# TYPE YOUR RUN COMMAND HERE
+"""
+
+
+def main_local(hparams):
+    main(hparams, None, None)
+
+
+def main(hparams, cluster, results_dict):
+    """
+    Main training routine specific for this project
+    :param hparams:
+    :return:
+    """
+    # ------------------------
+    # 1 INIT LIGHTNING MODEL
+    # ------------------------
+    print('loading model...')
+    model = LightningTemplateModel(hparams)
+    print('model built')
+
+    # ------------------------
+    # 2 INIT TEST TUBE EXP
+    # ------------------------
+    # when using grid search, it's possible for all models to start at once
+    # and use the same test tube experiment version
+    relative_node_id = int(os.environ['SLURM_NODEID'])
+    sleep(relative_node_id + 1)
+
+    # init experiment
+    exp = Experiment(
+        name=hyperparams.experiment_name,
+        save_dir=hyperparams.test_tube_save_path,
+        autosave=False,
+        description='test demo'
+    )
+
+    exp.argparse(hparams)
+    exp.save()
+
+    # ------------------------
+    # 3 DEFINE CALLBACKS
+    # ------------------------
+    model_save_path = '{}/{}/{}'.format(hparams.model_save_path, exp.name, exp.version)
+    early_stop = EarlyStopping(
+        monitor='val_acc',
+        patience=3,
+        verbose=True,
+        mode='max'
+    )
+
+    checkpoint = ModelCheckpoint(
+        filepath=model_save_path,
+        save_best_only=True,
+        verbose=True,
+        monitor='val_loss',
+        mode='min'
+    )
+
+    # ------------------------
+    # 4 INIT TRAINER
+    # ------------------------
+    trainer = Trainer(
+        experiment=exp,
+        cluster=cluster,
+        checkpoint_callback=checkpoint,
+        early_stop_callback=early_stop,
+        gpus=hparams.gpus,
+        nb_gpu_nodes=hyperparams.nb_gpu_nodes
+    )
+
+    # ------------------------
+    # 5 START TRAINING
+    # ------------------------
+    trainer.fit(model)
+
+
+def optimize_on_cluster(hyperparams):
+    # enable cluster training
+    # log all scripts to the test tube folder
+    cluster = SlurmCluster(
+        hyperparam_optimizer=hyperparams,
+        log_path=hyperparams.slurm_log_path,
+    )
+
+    # 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.per_experiment_nb_nodes = hyperparams.nb_gpu_nodes
+    cluster.job_time = '2:00:00'
+    cluster.gpu_type = 'volta'
+    cluster.memory_mb_per_node = 0
+
+    # any modules for code to run in env
+    cluster.add_command('source activate lightning')
+
+    # run only on 32GB voltas
+    cluster.add_slurm_cmd(cmd='constraint', value='volta32gb', comment='use 32gb gpus')
+    cluster.add_slurm_cmd(cmd='partition', value=hyperparams.gpu_partition, comment='use 32gb gpus')
+
+    # run hopt
+    # creates and submits jobs to slurm
+    cluster.optimize_parallel_cluster_gpu(
+        main,
+        nb_trials=hyperparams.nb_hopt_trials,
+        job_name=hyperparams.experiment_name
+    )
+
+
+if __name__ == '__main__':
+
+    # use default args
+    root_dir = os.path.dirname(os.path.realpath(__file__))
+    demo_log_dir = os.path.join(root_dir, 'pt_lightning_demo_logs')
+
+    checkpoint_dir = os.path.join(demo_log_dir, 'model_weights')
+    test_tube_dir = os.path.join(demo_log_dir, 'test_tube_data')
+    slurm_out_dir = os.path.join(demo_log_dir, 'slurm_scripts')
+
+    parent_parser = HyperOptArgumentParser(strategy='grid_search', add_help=False)
+
+    # cluster args not defined inside the model
+    parent_parser.add_argument('--gpu_partition', type=str, help='consult your cluster manual')
+
+    # TODO: make 1 param
+    parent_parser.add_argument('--per_experiment_nb_gpus', type=int, help='how many gpus to use in a node')
+    parent_parser.add_argument('--gpus', type=str, default='-1', help='how many gpus to use in the node')
+
+    parent_parser.add_argument('--nb_gpu_nodes', type=int, default=1, help='how many nodes to use in a cluster')
+    parent_parser.add_argument('--test_tube_save_path', type=str, default=test_tube_dir, help='where to save logs')
+    parent_parser.add_argument('--slurm_log_path', type=str, default=slurm_out_dir, help='where to save slurm meta')
+    parent_parser.add_argument('--model_save_path', type=str, default=checkpoint_dir, help='where to save model')
+    parent_parser.add_argument('--experiment_name', type=str, default='pt_lightning_exp_a', help='test tube exp name')
+    parent_parser.add_argument('--nb_hopt_trials', type=int, default=1, help='how many grid search trials to run')
+
+    # allow model to overwrite or extend args
+    parser = LightningTemplateModel.add_model_specific_args(parent_parser, root_dir)
+    hyperparams = parser.parse_args()
+
+    # ---------------------
+    # RUN TRAINING
+    # ---------------------
+    # run on HPC cluster
+    print('RUNNING ON SLURM CLUSTER')
+    optimize_on_cluster(hyperparams)

examples/new_project_templates/single_cpu_template.py

@@ -0,0 +1,110 @@
+"""
+Runs a model on a single node across N-gpus.
+"""
+import os
+import sys
+import numpy as np
+from time import sleep
+import torch
+
+from test_tube import HyperOptArgumentParser, Experiment, SlurmCluster
+from pytorch_lightning.models.trainer import Trainer
+from pytorch_lightning.utils.arg_parse import add_default_args
+
+from pytorch_lightning.callbacks import EarlyStopping, ModelCheckpoint
+
+SEED = 2334
+torch.manual_seed(SEED)
+np.random.seed(SEED)
+
+from lightning_module_template import LightningTemplateModel
+
+
+def main(hparams):
+    """
+    Main training routine specific for this project
+    :param hparams:
+    :return:
+    """
+    # ------------------------
+    # 1 INIT LIGHTNING MODEL
+    # ------------------------
+    print('loading model...')
+    model = LightningTemplateModel(hparams)
+    print('model built')
+
+    # ------------------------
+    # 2 INIT TEST TUBE EXP
+    # ------------------------
+
+    # init experiment
+    exp = Experiment(
+        name=hyperparams.experiment_name,
+        save_dir=hyperparams.test_tube_save_path,
+        autosave=False,
+        description='test demo'
+    )
+
+    exp.argparse(hparams)
+    exp.save()
+
+    # ------------------------
+    # 3 DEFINE CALLBACKS
+    # ------------------------
+    model_save_path = '{}/{}/{}'.format(hparams.model_save_path, exp.name, exp.version)
+    early_stop = EarlyStopping(
+        monitor='val_acc',
+        patience=3,
+        verbose=True,
+        mode='max'
+    )
+
+    checkpoint = ModelCheckpoint(
+        filepath=model_save_path,
+        save_best_only=True,
+        verbose=True,
+        monitor='val_loss',
+        mode='min'
+    )
+
+    # ------------------------
+    # 4 INIT TRAINER
+    # ------------------------
+    trainer = Trainer(
+        experiment=exp,
+        checkpoint_callback=checkpoint,
+        early_stop_callback=early_stop,
+    )
+
+    # ------------------------
+    # 5 START TRAINING
+    # ------------------------
+    trainer.fit(model)
+
+
+if __name__ == '__main__':
+
+    # dirs
+    root_dir = os.path.dirname(os.path.realpath(__file__))
+    demo_log_dir = os.path.join(root_dir, 'pt_lightning_demo_logs')
+    checkpoint_dir = os.path.join(demo_log_dir, 'model_weights')
+    test_tube_dir = os.path.join(demo_log_dir, 'test_tube_data')
+
+    # although we user hyperOptParser, we are using it only as argparse right now
+    parent_parser = HyperOptArgumentParser(strategy='grid_search', add_help=False)
+
+    # gpu args
+    parent_parser.add_argument('--test_tube_save_path', type=str, default=test_tube_dir, help='where to save logs')
+    parent_parser.add_argument('--model_save_path', type=str, default=checkpoint_dir, help='where to save model')
+    parent_parser.add_argument('--experiment_name', type=str, default='pt_lightning_exp_a', help='test tube exp name')
+
+    # allow model to overwrite or extend args
+    parser = LightningTemplateModel.add_model_specific_args(parent_parser, root_dir)
+    hyperparams = parser.parse_args()
+
+    # ---------------------
+    # RUN TRAINING
+    # ---------------------
+    # run on HPC cluster
+    print(f'RUNNING ON CPU')
+    main(hyperparams)

examples/new_project_templates/single_gpu_node_16bit_template.py

@@ -0,0 +1,113 @@
+"""
+Runs a model on a single node across N-gpus.
+"""
+import os
+import sys
+import numpy as np
+from time import sleep
+import torch
+
+from test_tube import HyperOptArgumentParser, Experiment, SlurmCluster
+from pytorch_lightning.models.trainer import Trainer
+from pytorch_lightning.utils.arg_parse import add_default_args
+
+from pytorch_lightning.callbacks import EarlyStopping, ModelCheckpoint
+
+SEED = 2334
+torch.manual_seed(SEED)
+np.random.seed(SEED)
+
+from lightning_module_template import LightningTemplateModel
+
+
+def main(hparams):
+    """
+    Main training routine specific for this project
+    :param hparams:
+    :return:
+    """
+    # ------------------------
+    # 1 INIT LIGHTNING MODEL
+    # ------------------------
+    print('loading model...')
+    model = LightningTemplateModel(hparams)
+    print('model built')
+
+    # ------------------------
+    # 2 INIT TEST TUBE EXP
+    # ------------------------
+
+    # init experiment
+    exp = Experiment(
+        name=hyperparams.experiment_name,
+        save_dir=hyperparams.test_tube_save_path,
+        autosave=False,
+        description='test demo'
+    )
+
+    exp.argparse(hparams)
+    exp.save()
+
+    # ------------------------
+    # 3 DEFINE CALLBACKS
+    # ------------------------
+    model_save_path = '{}/{}/{}'.format(hparams.model_save_path, exp.name, exp.version)
+    early_stop = EarlyStopping(
+        monitor='val_acc',
+        patience=3,
+        verbose=True,
+        mode='max'
+    )
+
+    checkpoint = ModelCheckpoint(
+        filepath=model_save_path,
+        save_best_only=True,
+        verbose=True,
+        monitor='val_loss',
+        mode='min'
+    )
+
+    # ------------------------
+    # 4 INIT TRAINER
+    # ------------------------
+    trainer = Trainer(
+        experiment=exp,
+        checkpoint_callback=checkpoint,
+        early_stop_callback=early_stop,
+        gpus=hparams.gpus,
+        use_amp=True
+    )
+
+    # ------------------------
+    # 5 START TRAINING
+    # ------------------------
+    trainer.fit(model)
+
+
+if __name__ == '__main__':
+
+    # dirs
+    root_dir = os.path.dirname(os.path.realpath(__file__))
+    demo_log_dir = os.path.join(root_dir, 'pt_lightning_demo_logs')
+    checkpoint_dir = os.path.join(demo_log_dir, 'model_weights')
+    test_tube_dir = os.path.join(demo_log_dir, 'test_tube_data')
+
+    # although we user hyperOptParser, we are using it only as argparse right now
+    parent_parser = HyperOptArgumentParser(strategy='grid_search', add_help=False)
+
+    # gpu args
+    parent_parser.add_argument('--gpus', type=str, default='-1', help='how many gpus to use in the node. -1 uses all the gpus on the node')
+    parent_parser.add_argument('--test_tube_save_path', type=str, default=test_tube_dir, help='where to save logs')
+    parent_parser.add_argument('--model_save_path', type=str, default=checkpoint_dir, help='where to save model')
+    parent_parser.add_argument('--experiment_name', type=str, default='pt_lightning_exp_a', help='test tube exp name')
+
+    # allow model to overwrite or extend args
+    parser = LightningTemplateModel.add_model_specific_args(parent_parser, root_dir)
+    hyperparams = parser.parse_args()
+
+    # ---------------------
+    # RUN TRAINING
+    # ---------------------
+    # run on HPC cluster
+    print(f'RUNNING INTERACTIVE MODE ON GPUS. gpu ids: {hyperparams.gpus}')
+    main(hyperparams)

examples/new_project_templates/single_gpu_node_template.py

@@ -0,0 +1,112 @@
+"""
+Runs a model on a single node across N-gpus.
+"""
+import os
+import sys
+import numpy as np
+from time import sleep
+import torch
+
+from test_tube import HyperOptArgumentParser, Experiment, SlurmCluster
+from pytorch_lightning.models.trainer import Trainer
+from pytorch_lightning.utils.arg_parse import add_default_args
+
+from pytorch_lightning.callbacks import EarlyStopping, ModelCheckpoint
+
+SEED = 2334
+torch.manual_seed(SEED)
+np.random.seed(SEED)
+
+from lightning_module_template import LightningTemplateModel
+
+
+def main(hparams):
+    """
+    Main training routine specific for this project
+    :param hparams:
+    :return:
+    """
+    # ------------------------
+    # 1 INIT LIGHTNING MODEL
+    # ------------------------
+    print('loading model...')
+    model = LightningTemplateModel(hparams)
+    print('model built')
+
+    # ------------------------
+    # 2 INIT TEST TUBE EXP
+    # ------------------------
+
+    # init experiment
+    exp = Experiment(
+        name=hyperparams.experiment_name,
+        save_dir=hyperparams.test_tube_save_path,
+        autosave=False,
+        description='test demo'
+    )
+
+    exp.argparse(hparams)
+    exp.save()
+
+    # ------------------------
+    # 3 DEFINE CALLBACKS
+    # ------------------------
+    model_save_path = '{}/{}/{}'.format(hparams.model_save_path, exp.name, exp.version)
+    early_stop = EarlyStopping(
+        monitor='val_acc',
+        patience=3,
+        verbose=True,
+        mode='max'
+    )
+
+    checkpoint = ModelCheckpoint(
+        filepath=model_save_path,
+        save_best_only=True,
+        verbose=True,
+        monitor='val_loss',
+        mode='min'
+    )
+
+    # ------------------------
+    # 4 INIT TRAINER
+    # ------------------------
+    trainer = Trainer(
+        experiment=exp,
+        checkpoint_callback=checkpoint,
+        early_stop_callback=early_stop,
+        gpus=hparams.gpus,
+    )
+
+    # ------------------------
+    # 5 START TRAINING
+    # ------------------------
+    trainer.fit(model)
+
+
+if __name__ == '__main__':
+
+    # dirs
+    root_dir = os.path.dirname(os.path.realpath(__file__))
+    demo_log_dir = os.path.join(root_dir, 'pt_lightning_demo_logs')
+    checkpoint_dir = os.path.join(demo_log_dir, 'model_weights')
+    test_tube_dir = os.path.join(demo_log_dir, 'test_tube_data')
+
+    # although we user hyperOptParser, we are using it only as argparse right now
+    parent_parser = HyperOptArgumentParser(strategy='grid_search', add_help=False)
+
+    # gpu args
+    parent_parser.add_argument('--gpus', type=str, default='-1', help='how many gpus to use in the node. -1 uses all the gpus on the node')
+    parent_parser.add_argument('--test_tube_save_path', type=str, default=test_tube_dir, help='where to save logs')
+    parent_parser.add_argument('--model_save_path', type=str, default=checkpoint_dir, help='where to save model')
+    parent_parser.add_argument('--experiment_name', type=str, default='pt_lightning_exp_a', help='test tube exp name')
+
+    # allow model to overwrite or extend args
+    parser = LightningTemplateModel.add_model_specific_args(parent_parser, root_dir)
+    hyperparams = parser.parse_args()
+
+    # ---------------------
+    # RUN TRAINING
+    # ---------------------
+    # run on HPC cluster
+    print(f'RUNNING INTERACTIVE MODE ON GPUS. gpu ids: {hyperparams.gpus}')
+    main(hyperparams)

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.utils.arg_parse import add_default_args
-from pytorch-lightning.utils.pt_callbacks import EarlyStopping, ModelCheckpoint
-from demo.example_model import ExampleModel
+from pytorch_lightning.models.trainer import Trainer
+from pytorch_lightning.utils.arg_parse import add_default_args
+from pytorch_lightning.callbacks.pt_callbacks import EarlyStopping, ModelCheckpoint
+from docs.source.examples.example_model import ExampleModel
 
 
 def main(hparams):
@@ -41,7 +41,6 @@ def main(hparams):
     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='val_acc',

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,2 @@
+from .models import Trainer
+from .root_module.root_module import LightningModule

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 LightningDistributedDataParallel
 
 
 class Callback(object):
@@ -33,6 +34,8 @@ class Callback(object):
         self.params = params
 
     def set_model(self, model):
+        if type(model) is LightningDistributedDataParallel:
+            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

@@ -1,11 +1,24 @@
-import torch
-import tqdm
-import numpy as np
-from pytorch_lightning.root_module.memory import get_gpu_memory_map
+"""
+The trainer handles all the logic for running a val loop, training loop, distributing, etc...
+"""
+from time import sleep
+import subprocess
 import traceback
-from pytorch_lightning.root_module.model_saving import TrainerIO
+import warnings
+import os
+
+import torch
+from torch.utils.data.distributed import DistributedSampler
 from torch.optim.lr_scheduler import MultiStepLR
-import pdb
+import torch.multiprocessing as mp
+import torch.distributed as dist
+import numpy as np
+import tqdm
+
+from pytorch_lightning.root_module.memory import get_gpu_memory_map
+from pytorch_lightning.root_module.model_saving import TrainerIO
+from pytorch_lightning.pt_overrides.override_data_parallel import LightningDistributedDataParallel
+
 
 try:
     from apex import amp
@@ -13,35 +26,42 @@ try:
 except ModuleNotFoundError:
     APEX_AVAILABLE = False
 
+
 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,
-                 enable_tqdm=True,
+                 nb_gpu_nodes=1,
+                 gpus=None,
+                 progress_bar=True,
                  overfit_pct=0.0,
                  track_grad_norm=-1,
                  check_val_every_n_epoch=1,
                  fast_dev_run=False,
                  accumulate_grad_batches=1,
-                 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,
+                 print_nan_grads=False,
+                 amp_level='O2',
                  nb_sanity_val_steps=5):
 
         # Transfer params
+        self.nb_gpu_nodes = nb_gpu_nodes
+        self.gradient_clip = gradient_clip
         self.check_val_every_n_epoch = check_val_every_n_epoch
         self.enable_early_stop = enable_early_stop
         self.track_grad_norm = track_grad_norm
         self.fast_dev_run = fast_dev_run
-        self.on_gpu = on_gpu
-        self.enable_tqdm = enable_tqdm
+        self.on_gpu = gpus is not None and torch.cuda.is_available()
+        self.progress_bar = progress_bar
         self.experiment = experiment
         self.exp_save_path = experiment.get_data_path(experiment.name, experiment.version)
         self.cluster = cluster
@@ -58,6 +78,28 @@ class Trainer(TrainerIO):
         self.nb_sanity_val_steps = nb_sanity_val_steps
         self.lr_scheduler_milestones = [] if lr_scheduler_milestones is None else [int(x.strip()) for x in lr_scheduler_milestones.split(',')]
         self.lr_schedulers = []
+        self.amp_level = amp_level
+        self.print_nan_grads = print_nan_grads
+        self.data_parallel_device_ids = None
+        self.world_size = 1
+
+        # gpus come in as a string.
+        # if gpus = -1 then use all available devices
+        # otherwise, split the string using commas
+        if gpus is not None:
+            if gpus == '-1':
+                self.data_parallel_device_ids = list(range(0, torch.cuda.device_count()))
+            else:
+                self.data_parallel_device_ids = [int(x.strip()) for x in gpus.split(',')]
+
+            # set the correct cuda visible devices (using pci order)
+            os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
+            os.environ["CUDA_VISIBLE_DEVICES"] = ','.join([str(x) for x in self.data_parallel_device_ids])
+
+        self.data_parallel = self.data_parallel_device_ids is not None and len(self.data_parallel_device_ids) > 0
+
+        # process info
+        self.proc_rank = 0
 
         # training state
         self.optimizers = None
@@ -80,11 +122,21 @@ class Trainer(TrainerIO):
         self.__determine_data_use_amount(train_percent_check, val_percent_check, test_percent_check, overfit_pct)
         print('gpu available: {}, used: {}'.format(torch.cuda.is_available(), self.on_gpu))
 
-        # apex test
+        # 16 bit mixed precision training using apex
         self.use_amp = use_amp and APEX_AVAILABLE
         if self.use_amp:
             print('using 16bit precision')
 
+        if use_amp and not APEX_AVAILABLE:
+            msg = '''
+            You set use_amp=True but do not have apex installed.
+            Install apex first using this guide and rerun with use_amp=True: 
+            https://github.com/NVIDIA/apex#linux
+            
+            this run will NOT use 16 bit precision
+            '''
+            warnings.warn(msg)
+
     def __determine_data_use_amount(self, train_percent_check, val_percent_check, test_percent_check, overfit_pct):
         """
         Use less data for debugging purposes
@@ -105,12 +157,15 @@ class Trainer(TrainerIO):
     def __tng_tqdm_dic(self):
         tqdm_dic = {
             'tng_loss': '{0:.3f}'.format(self.avg_loss),
-            'gpu': '{}'.format(self.current_gpu_name),
             'v_nb': '{}'.format(self.experiment.version),
             'epoch': '{}'.format(self.current_epoch),
             'batch_nb':'{}'.format(self.batch_nb),
         }
         tqdm_dic.update(self.tqdm_metrics)
+
+        if self.on_gpu:
+            tqdm_dic['gpu'] = '{}'.format(self.current_gpu_name)
+
         return tqdm_dic
 
     def __layout_bookeeping(self):
@@ -122,24 +177,27 @@ class Trainer(TrainerIO):
         self.tqdm_metrics = {}
 
         # determine number of training batches
-        nb_tng_batches = self.model.nb_batches(self.tng_dataloader)
-        self.nb_tng_batches = int(nb_tng_batches * self.train_percent_check)
+        self.nb_tng_batches = len(self.tng_dataloader)
+        self.nb_tng_batches = int(self.nb_tng_batches * self.train_percent_check)
 
         # determine number of validation batches
-        nb_val_batches = self.model.nb_batches(self.val_dataloader)
-        nb_val_batches = int(nb_val_batches * self.val_percent_check)
-        nb_val_batches = max(1, nb_val_batches)
-        self.nb_val_batches = nb_val_batches
+        self.nb_val_batches = len(self.val_dataloader)
+        self.nb_val_batches = int(self.nb_val_batches * self.val_percent_check)
+        self.nb_val_batches = max(1, self.nb_val_batches)
+        self.nb_val_batches = self.nb_val_batches
 
         # determine number of test batches
-        nb_test_batches = self.model.nb_batches(self.test_dataloader)
-        self.nb_test_batches = int(nb_test_batches * self.test_percent_check)
+        self.nb_test_batches = len(self.test_dataloader)
+        self.nb_test_batches = int(self.nb_test_batches * self.test_percent_check)
 
         # determine when to check validation
-        self.val_check_batch = int(nb_tng_batches * self.val_check_interval)
+        self.val_check_batch = int(self.nb_tng_batches * self.val_check_interval)
 
     def __add_tqdm_metrics(self, metrics):
         for k, v in metrics.items():
+            if type(v) is torch.Tensor:
+                v = v.item()
+
             self.tqdm_metrics[k] = v
 
     def validate(self, model, dataloader, max_batches):
@@ -150,7 +208,8 @@ class Trainer(TrainerIO):
         :param max_batches: Scalar
         :return:
         """
-        print('validating...')
+        if self.proc_rank == 0:
+            print('validating...')
 
         # enable eval mode
         model.zero_grad()
@@ -163,33 +222,41 @@ class Trainer(TrainerIO):
         outputs = []
 
         # run training
-        for i, data_batch in enumerate(dataloader):
+        for batch_i, data_batch in enumerate(dataloader):
 
             if data_batch is None:
                 continue
 
             # stop short when on fast dev run
-            if max_batches is not None and i >= max_batches:
+            if max_batches is not None and batch_i >= max_batches:
                 break
 
             # -----------------
             # RUN VALIDATION STEP
             # -----------------
-            output = model.validation_step(data_batch)
+            if self.data_parallel:
+                output = model(data_batch, batch_i)
+            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
-        val_results = model.validation_end(outputs)
+        if self.data_parallel:
+            val_results = model.module.validation_end(outputs)
+        else:
+            val_results = model.validation_end(outputs)
 
         # enable train mode again
         model.train()
 
         # enable gradients to save memory
         torch.set_grad_enabled(True)
+
         return val_results
 
     def __get_dataloaders(self, model):
@@ -202,30 +269,174 @@ class Trainer(TrainerIO):
         self.test_dataloader = model.test_dataloader
         self.val_dataloader = model.val_dataloader
 
+        if self.on_gpu and type(self.tng_dataloader.sampler) is not DistributedSampler:
+            msg = '''
+            when using multiple gpus and multiple nodes you must pass a DistributedSampler to DataLoader(sampler).
+            
+            ie: this:
+            dataset = myDataset()
+            dataloader = Dataloader(dataset)
+            
+            becomes:
+            dataset = myDataset()
+            dist_sampler = torch.utils.data.distributed.DistributedSampler(dataset)
+            dataloader = Dataloader(dataset, sampler=dist_sampler)
+            '''
+            raise Exception(msg)
+
     # -----------------------------
     # MODEL TRAINING
     # -----------------------------
     def fit(self, model):
-        self.model = model
-        model.trainer = self
 
-        # transfer data loaders from model
-        self.__get_dataloaders(model)
+        # when using gpus, first thing we do is spawn a new process between each worker
+        # applies to single gpu, multi-gpu and multi-nodes
+        if self.on_gpu:
+            self.experiment = self.experiment.get_meta_copy()
+            mp.spawn(self.dp_train, nprocs=len(self.data_parallel_device_ids), args=(model, ))
+        else:
+            # CHOOSE OPTIMIZER
+            # filter out the weights that were done on gpu so we can load on good old cpus
+            self.optimizers = model.configure_optimizers()
 
-        # init training constants
-        self.__layout_bookeeping()
+            # run through amp wrapper
+            if self.use_amp:
+                # An example
+                model, optimizers = amp.initialize(
+                    model, self.optimizers, opt_level=self.amp_level,
+                )
+                self.optimizers = optimizers
+
+            self.__run_pretrain_routine(model)
+
+    def dp_train(self, gpu_nb, model):
+        """
+        Entry point into a DP thread
+        :param gpu_nb:
+        :param model:
+        :param cluster_obj:
+        :return:
+        """
+        # node rank using relative slurm id
+        # otherwise default to node rank 0
+        try:
+            node_rank = int(os.environ['SLURM_NODEID'])
+        except KeyError as e:
+            node_rank = 0
+
+        # recover original exp before went into process
+        self.experiment = self.experiment.get_non_ddp_exp()
+
+        # show progbar only on prog_rank 0
+        self.prog_bar = self.prog_bar and node_rank == 0 and gpu_nb == 0
+
+        # determine which process we are and world size
+        self.proc_rank = node_rank * len(self.data_parallel_device_ids) + gpu_nb
+        self.world_size = self.nb_gpu_nodes * len(self.data_parallel_device_ids)
+
+        # set up server using proc 0's ip address
+        # try to init for 20 times at max in case ports are taken
+        ip = self.__get_root_node_ip(self.proc_rank, self.nb_gpu_nodes)
+        self.__init_tcp_connection(ip)
 
         # CHOOSE OPTIMIZER
         # filter out the weights that were done on gpu so we can load on good old cpus
         self.optimizers = model.configure_optimizers()
 
+        # MODEL
+        # copy model to each gpu
+        torch.cuda.set_device(gpu_nb)
+        model.cuda(gpu_nb)
+
+        # AMP
+        # run through amp wrapper before going to distributed DP
         if self.use_amp:
             # An example
-            self.model, optimizer = amp.initialize(
-                self.model, self.optimizers[0], opt_level="O2",
+            model, optimizers = amp.initialize(
+                model, self.optimizers, opt_level=self.amp_level,
             )
-            self.optimizers[0] = optimizer
-            model.trainer = self
+            self.optimizers = optimizers
+
+        model = LightningDistributedDataParallel(model, device_ids=[gpu_nb])
+
+        # continue training routine
+        self.__run_pretrain_routine(model)
+
+    def __init_tcp_connection(self, ip, port=12000, tries=0):
+        if tries > 20:
+            raise RuntimeError('Failed to connect using 20 different ip addresses')
+
+        try:
+            dist.init_process_group("nccl", init_method=f'tcp://{ip}:{port}', rank=self.proc_rank, world_size=self.world_size)
+        except RuntimeError as e:
+            # port taken
+            warnings.warn(f'port {port} taken, trying port {port}...')
+            self.__init_tcp_connection(ip, port + 1, tries + 1)
+
+    def __get_root_node_ip(self, world_gpu_nb, nb_gpu_nodes):
+        """
+        Resolves the ip address of proc 0.
+        Proc 0 writes address to a file. Every other process waits until the ip is available before it starts
+
+        :param world_gpu_nb: gpu number amongst all the world gpus
+        :param nb_gpu_nodes:
+        :param ip_file_dir:
+        :return:
+        """
+        # on one node we use localhost
+        if nb_gpu_nodes == 1:
+            return '127.0.0.1'
+
+        # where to store ip_table
+        ip_file_dir = os.path.join(self.cluster.log_path, 'ip_tables')
+
+        # the first gpu in the world becomes the host
+        # this is based on its global rank
+        # it communicates its ip by saving an ip_table to the slurm cluster logging dir
+        # every other process waits for this ip to appear before continuing
+        ip_table_name = f'.ip_meta_' + os.environ['SLURM_JOB_ID']
+        ip_file = os.path.join(ip_file_dir, ip_table_name)
+        os.makedirs(ip_file_dir, exist_ok=True)
+
+        if world_gpu_nb == 0:
+            # get the proc 0 IP
+            root_ip = subprocess.run(['hostname', '-I'], stdout=subprocess.PIPE).stdout.decode('utf-8')
+            root_ip = root_ip.split(' ')[0]
+
+            # save the ip to the file
+            with open(file=ip_file, mode='w') as f:
+                f.write(root_ip)
+
+            return root_ip
+        else:
+            # wait up to 120 seconds until proc 0 writes
+            # once written, read proc 0's address and use it to configure server
+            for i in range(0, 120):
+                sleep(1.0)
+                if os.path.exists(ip_file):
+                    ip = list(open(file=ip_file, mode='r'))[0]
+                    return ip
+
+    def __run_pretrain_routine(self, model):
+        """
+        Sanity check a few things before starting actual training
+        :param model:
+        :return:
+        """
+        ref_model = model
+        if self.on_gpu:
+            ref_model = model.module
+
+        ref_model.trainer = self
+
+        # set local properties on the model
+        ref_model.on_gpu = self.on_gpu
+
+        # transfer data loaders from model
+        self.__get_dataloaders(ref_model)
+
+        # init training constants
+        self.__layout_bookeeping()
 
         # add lr schedulers
         if self.lr_scheduler_milestones is not None:
@@ -234,17 +445,19 @@ class Trainer(TrainerIO):
                 self.lr_schedulers.append(scheduler)
 
         # print model summary
-        model.summarize()
+        if self.proc_rank == 0:
+            ref_model.summarize()
 
-        # put on gpu if needed
-        if self.on_gpu:
-            model = model.cuda()
+        # give model convenience properties
+        ref_model.trainer = self
+        ref_model.experiment = self.experiment
 
         # run tiny validation to make sure program won't crash during val
         _ = self.validate(model, self.val_dataloader, max_batches=self.nb_sanity_val_steps)
 
         # save exp to get started
-        self.experiment.save()
+        if self.proc_rank == 0:
+            self.experiment.save()
 
         # enable cluster checkpointing
         if self.cluster is not None:
@@ -253,6 +466,7 @@ class Trainer(TrainerIO):
         # ---------------------------
         # CORE TRAINING LOOP
         # ---------------------------
+        self.model = model
         self.__train()
 
     def __train(self):
@@ -262,24 +476,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
@@ -290,7 +508,7 @@ class Trainer(TrainerIO):
                 # ---------------
                 # RUN TRAIN STEP
                 # ---------------
-                batch_result = self.__run_tng_batch(data_batch)
+                batch_result = self.__run_tng_batch(data_batch, batch_nb)
                 early_stop_epoch = batch_result == -1
 
                 # ---------------
@@ -302,14 +520,18 @@ class Trainer(TrainerIO):
 
                 # when batch should be saved
                 if (batch_nb + 1) % self.log_save_interval == 0 or early_stop_epoch:
-                    self.experiment.save()
+                    if self.proc_rank == 0:
+                        self.experiment.save()
 
                 # when metrics should be logged
                 if batch_nb % self.add_log_row_interval == 0 or early_stop_epoch:
                     # count items in memory
                     # nb_params, nb_tensors = count_mem_items()
 
-                    metrics = self.model.update_tng_log_metrics(self.__tng_tqdm_dic)
+                    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
                     if self.on_gpu:
@@ -318,16 +540,21 @@ 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
-                    self.experiment.log(metrics)
-                    self.experiment.save()
+                    scalar_metrics = self.__metrics_to_scalars(metrics, blacklist=self.__log_vals_blacklist())
+                    if self.proc_rank == 0:
+                        self.experiment.log(scalar_metrics, global_step=self.global_step)
+                        self.experiment.save()
 
                 # hook
                 if self.__is_function_implemented('on_batch_end'):
-                    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:
@@ -335,7 +562,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:
@@ -347,38 +575,87 @@ class Trainer(TrainerIO):
                 if stop:
                     return
 
+    def __metrics_to_scalars(self, metrics, blacklist=[]):
+        new_metrics = {}
+        for k, v in metrics.items():
+            if type(v) is torch.Tensor:
+                v = v.item()
 
-    def __run_tng_batch(self, data_batch):
+            if type(v) is dict:
+                v = self.__metrics_to_scalars(v)
+
+            if k not in blacklist:
+                new_metrics[k] = float(v)
+
+        return new_metrics
+
+    def __log_vals_blacklist(self):
+        """avoid logging some vals lightning uses to maintain state"""
+        blacklist = {'batch_nb', 'v_nb', 'epoch', 'gpu'}
+        return blacklist
+
+    def __run_tng_batch(self, data_batch, batch_nb):
         if data_batch is None:
             return 0
 
         # hook
         if self.__is_function_implemented('on_batch_start'):
-            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)
+        if self.data_parallel:
+            output = self.model(data_batch, batch_nb)
+        else:
+            output = self.model.training_step(data_batch, batch_nb)
+
+        try:
+            model_specific_tqdm_metrics_dic = output['tqdm_metrics']
+        except Exception as e:
+            model_specific_tqdm_metrics_dic = {}
+
+        # if output dict doesn't have the keyword loss
+        # then assume the output=loss if scalar
+        try:
+            loss = output['loss']
+        except Exception as e:
+            if type(loss) is torch.Tensor:
+                loss = output
+
         self.__add_tqdm_metrics(model_specific_tqdm_metrics_dic)
 
         # backward pass
         if self.use_amp:
+            # scale loss when using amp
             for optimizer in self.optimizers:
                 with amp.scale_loss(loss, optimizer) as scaled_loss:
                     scaled_loss.backward()
         else:
             loss.backward()
 
+        if self.print_nan_grads:
+            model = self.model.module if self.data_parallel else self.model
+            for param in model.parameters():
+                print(param.grad.float().sum())
+
+        # avoid memory leaks
         self.batch_loss_value += loss.item()
 
         # gradient update with accumulated gradients
         if (self.batch_nb + 1) % self.accumulate_grad_batches == 0:
 
+            # clip gradients
+            if self.gradient_clip > 0:
+                model = self.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()
@@ -395,7 +672,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)
@@ -437,11 +714,12 @@ 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)
 
         # model checkpointing
-        print('save callback...')
-        self.checkpoint_callback.on_epoch_end(epoch=self.current_epoch, logs=self.__tng_tqdm_dic)
+        if self.proc_rank == 0:
+            print('save callback...')
+            self.checkpoint_callback.on_epoch_end(epoch=self.current_epoch, logs=self.__tng_tqdm_dic)

pytorch_lightning/pt_overrides/override_data_parallel.py

@@ -0,0 +1,164 @@
+from torch.nn import DataParallel
+from torch.nn.parallel import DistributedDataParallel
+import itertools
+
+import threading
+import torch
+from torch.cuda._utils import _get_device_index
+import pdb
+
+
+def _find_tensors(obj):
+    r"""
+    Recursively find all tensors contained in the specified object.
+    """
+    if isinstance(obj, torch.Tensor):
+        return [obj]
+    if isinstance(obj, (list, tuple)):
+        return itertools.chain(*map(_find_tensors, obj))
+    if isinstance(obj, dict):
+        return itertools.chain(*map(_find_tensors, obj.values()))
+    return []
+
+
+
+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)])
+
+
+class LightningDistributedDataParallel(DistributedDataParallel):
+    """
+    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 forward(self, *inputs, **kwargs):
+        self._sync_params()
+        if self.device_ids:
+            inputs, kwargs = self.scatter(inputs, kwargs, self.device_ids)
+            if len(self.device_ids) == 1:
+                # --------------
+                # LIGHTNING MOD
+                # --------------
+                # normal
+                # output = self.module(*inputs[0], **kwargs[0])
+
+                # lightning
+                if self.module.training:
+                    output = self.module.training_step(*inputs[0], **kwargs[0])
+                else:
+                    output = self.module.validation_step(*inputs[0], **kwargs[0])
+            else:
+                outputs = self.parallel_apply(self._module_copies[:len(inputs)], inputs, kwargs)
+                output = self.gather(outputs, self.output_device)
+        else:
+            output = self.module(*inputs, **kwargs)
+
+        if torch.is_grad_enabled():
+            # We'll return the output object verbatim since it is a freeform
+            # object. We need to find any tensors in this object, though,
+            # because we need to figure out which parameters were used during
+            # this forward pass, to ensure we short circuit reduction for any
+            # unused parameters. Only if `find_unused_parameters` is set.
+            if self.find_unused_parameters:
+                self.reducer.prepare_for_backward(list(_find_tensors(output)))
+            else:
+                self.reducer.prepare_for_backward([])
+        return output
+
+
+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 LightningDistributedDataParallel
 
 class ModelIO(object):
 
@@ -20,6 +21,23 @@ class ModelIO(object):
         """
         raise NotImplementedError
 
+    # -------------------------
+    # OPTIONAL HOOKS
+    # -------------------------
+    def on_hpc_save(self):
+        """
+        Hook to do whatever you need right before Slurm manager saves the model
+        :return:
+        """
+        pass
+
+    def on_hpc_load(self):
+        """
+        Hook to do whatever you need right before Slurm manager loads the model
+        :return:
+        """
+        pass
+
 
 class TrainerIO(object):
 
@@ -48,7 +66,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 LightningDistributedDataParallel else self.model
+        checkpoint_dict = model.get_save_dict()
 
         # merge trainer and model saving items
         checkpoint.update(checkpoint_dict)
@@ -99,15 +118,24 @@ 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()
 
+        # close experiment to avoid issues
+        experiment.close()
+
         ckpt_number = self.max_ckpt_in_folder(folderpath) + 1
 
         if not os.path.exists(folderpath):
             os.makedirs(folderpath, exist_ok=True)
         filepath = '{}/hpc_ckpt_{}.ckpt'.format(folderpath, ckpt_number)
 
+        # give model a chance to do something on hpc_save
+        self.on_hpc_save()
+
         # request what to save from the model
         checkpoint_dict = self.dump_checkpoint()
 
@@ -126,10 +154,18 @@ 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)
+
+        # call model hook
+        self.on_hpc_load()
 
     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,10 @@ 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
@@ -20,28 +20,17 @@ class RootModule(GradInformation, ModelIO, OptimizerConfig, ModelHooks):
         self.current_epoch = 0
         self.global_step = 0
         self.loaded_optimizer_states_dict = {}
-        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.experiment = None
 
         # track if gpu was requested for checkpointing
         self.on_gpu = False
-        try:
-            self.on_gpu = hparams.on_gpu
-        except Exception as e:
-            pass
 
         # computed vars for the dataloaders
         self._tng_dataloader = None
         self._val_dataloader = None
         self._test_dataloader = None
 
-        if self.on_gpu:
-            print('running on gpu...')
-            torch.set_default_tensor_type(hparams.default_tensor_type)
-
     def forward(self, *args, **kwargs):
         """
         Expand model in into whatever you need.
@@ -51,7 +40,7 @@ class RootModule(GradInformation, ModelIO, OptimizerConfig, ModelHooks):
         """
         raise NotImplementedError
 
-    def validation_step(self, data_batch):
+    def validation_step(self, data_batch, batch_nb):
         """
         return whatever outputs will need to be aggregated in validation_end
         :param data_batch:
@@ -67,7 +56,7 @@ class RootModule(GradInformation, ModelIO, OptimizerConfig, ModelHooks):
         """
         raise NotImplementedError
 
-    def training_step(self, data_batch):
+    def training_step(self, data_batch, batch_nb):
         """
         return loss, dict with metrics for tqdm
         :param data_batch:
@@ -103,9 +92,6 @@ class RootModule(GradInformation, ModelIO, OptimizerConfig, ModelHooks):
         model_summary = ModelSummary(self)
         print(model_summary)
 
-    def nb_batches(self, dataloader):
-        a = math.ceil(float(len(dataloader.dataset) / self.batch_size))
-        return int(a)
 
     def freeze(self):
         for param in self.parameters():

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
@@ -37,20 +39,21 @@ def add_default_args(parser, root_dir, rand_seed=None, possible_model_names=None
         parser.add_argument('--model_name', default='', help=','.join(possible_model_names))
 
     # test_tube settings
-    parser.add_argument('-en', '--tt_name', default='r_lib_')
-    parser.add_argument('-td', '--tt_description', default='test research  lib')
+    parser.add_argument('-en', '--tt_name', default='pt_test')
+    parser.add_argument('-td', '--tt_description', default='pytorch lightning test')
     parser.add_argument('--tt_save_path', default=root_dir + '/test_tube_logs', help='logging dir')
     parser.add_argument('--enable_single_run', dest='single_run', action='store_true')
     parser.add_argument('--nb_hopt_trials', default=1, type=int)
     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='0', type=str)
+    parser.add_argument('--gpus', default=None, 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')
+    parser.add_argument('--amp_level', default='O2',type=str)
+
 
     # run on hpc
     parser.add_argument('--on_cluster', dest='on_cluster', action='store_true')
@@ -65,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('--enable_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('--interactive', dest='interactive', action='store_true', help='runs on gpu without cluster')
+    parser.add_argument('--debug', dest='debug', action='store_true', help='enables/disables test tube')
+    parser.add_argument('--local', dest='local', action='store_true', help='enables local tng')
 
     # optimizer
     parser.add_argument('--lr_scheduler_milestones', default=None, type=str)

requirements.txt

@@ -1,8 +1,8 @@
+
 atomicwrites==1.2.1
 attrs==18.2.0
 certifi==2018.11.29
 cffi==1.11.5
-h5py==2.9.0
 imageio==2.4.1
 mkl-fft==1.0.6
 mkl-random==1.0.2
@@ -21,7 +21,15 @@ scikit-learn==0.20.2
 scipy==1.2.0
 six==1.12.0
 sklearn==0.0
-test-tube==0.6282
+tensorboard==1.14.0
+tensorboardX==1.7
+tensorflow==1.14.0
+test-tube==0.643
 torch==1.0.0
 torchvision==0.2.1
-tqdm==4.28.1
+tqdm==4.32.1
+twine==1.13.0
+urllib3==1.25.3
+webencodings==0.5.1
+Werkzeug==0.15.4
+wrapt==1.11.2

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.dev182',
+    version='0.2.2',
     description="The Keras for ML researchers using PyTorch",
     author="William Falcon",
     author_email="waf2107@columbia.edu",
@@ -17,9 +17,10 @@ setup(
     keywords=["deep learning", "pytorch", "AI"],
     python_requires=">=3.5",
     install_requires=[
-        "torch",
+        "torch>=1.1.0",
         "tqdm",
-        "test-tube",
+        "test-tube>=0.653",
+        "tensorflow>=1.14.0"
     ],
     packages=find_packages(),
     long_description=open("README.md", encoding="utf-8").read(),