release v0.2.5.1

add a hook for on_tng_metrics so that users get access to the grad_no…

.gitignore

@@ -0,0 +1,122 @@
+# project
+.DS_Store
+.data/
+run_configs/
+test_tube_logs/
+test_tube_data/
+datasets/
+model_weights/
+app/models/
+pip-wheel-metadata/
+test_tube_exp/
+
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+example.py
+timit_data/
+LJSpeech-1.1/
+
+# C extensions
+*.so
+
+.idea/
+
+# Distribution / packaging
+.Python
+env/
+ide_layouts/
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+.hypothesis/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# pyenv
+.python-version
+
+# celery beat schedule file
+celerybeat-schedule
+
+# SageMath parsed files
+*.sage.py
+
+# dotenv
+.env
+
+# virtualenv
+.venv
+venv/
+ENV/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+
+# data
+mnist/

MANIFEST.in

@@ -0,0 +1,9 @@
+graft docs
+
+include COPYING
+include AUTHORS
+
+recursive-include src/einsteinpy/tests *.py *.html
+
+prune docs/source/examples/.ipynb_checkpoints
+global-exclude *.py[cod] __pycache__ *.so *.dylib

README.md

@@ -1,177 +1,233 @@
-# Pytorch-lightning
-The Keras for ML-researchers in PyTorch.    
+<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/lightning_logo.png" width="50">
+  </a>
+</p>
+<h3 align="center">
+  Pytorch Lightning
+</h3>
+<p align="center">
+  The Keras for ML researchers using PyTorch. More control. Less boilerplate.    
+</p>
+<p align="center">
+  <a href="https://badge.fury.io/py/pytorch-lightning"><img src="https://badge.fury.io/py/pytorch-lightning.svg" alt="PyPI version" height="18"></a>
+<!--   <a href="https://travis-ci.org/williamFalcon/test-tube"><img src="https://travis-ci.org/williamFalcon/pytorch-lightning.svg?branch=master"></a> -->
+  <a href="https://github.com/williamFalcon/pytorch-lightning/blob/master/COPYING"><img src="https://img.shields.io/badge/License-MIT-yellow.svg"></a>
+</p>   
+
+```bash
+pip install pytorch-lightning    
+```
+
+## Docs   
+**[View the docs here](https://williamfalcon.github.io/pytorch-lightning/)**
+
+## What is it?  
+Keras and fast.ai are too abstract for researchers. Lightning abstracts the full training loop but gives you control in the critical points.   
+
+
+## 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.
 
-## Usage
 To use lightning do 2 things:  
-1. [Define a trainer](https://github.com/williamFalcon/pytorch-lightning/blob/master/pytorch_lightning/trainer_main.py) (which will run ALL your models).   
-2. [Define a model](https://github.com/williamFalcon/pytorch-lightning/blob/master/pytorch_lightning/models/sample_model_template/model_template.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).     
 
-### Example:   
+## What does lightning control for me?
+Everything!    
+Except for these 6 core functions which you define:    
 
-#### Define the trainer   
+```{.python}
+# what to do in the training loop
+def training_step(self, data_batch, batch_nb):
 
-```python
-# trainer.py
+# what to do in the validation loop
+def validation_step(self, data_batch, batch_nb):
 
-from pytorch_lightning.models.trainer import Trainer   
-from pytorch_lightning.utils.pt_callbacks import EarlyStopping, ModelCheckpoint
-from my_project import My_Model   
-from test_tube import HyperOptArgumentParser, Experiment, SlurmCluster
+# how to aggregate validation_step outputs
+def validation_end(self, outputs):
 
-# --------------
-# TEST TUBE INIT
-exp = Experiment(
-    name='my_exp',
-    debug=True,
-    save_dir='/some/path',
-    autosave=False,
-    description='my desc'
-)
-
-# --------------------
-# CALLBACKS
-early_stop = EarlyStopping(
-    monitor='val_loss',
-    patience=3,
-    verbose=True,
-    mode='min'
-)
-
-model_save_path = 'PATH/TO/SAVE'
-checkpoint = ModelCheckpoint(
-    filepath=model_save_path,
-    save_function=None,
-    save_best_only=True,
-    verbose=True,
-    monitor='val_acc',
-    mode='min'
-)
-
-# configure trainer
-trainer = Trainer(
-    experiment=experiment,
-    cluster=cluster,
-    checkpoint_callback=checkpoint,
-    early_stop_callback=early_stop
-)  
-
-# init model and train
-model = My_Model()
-trainer.fit(model)
+# and your dataloaders
+def tng_dataloader():
+def val_dataloader():
+def test_dataloader():
 ```
 
-#### Define the model   
+**Could be as complex as seq-2-seq + attention**    
 
 ```python
-from torch import nn
-
-class My_Model(RootModule):
-    def __init__(self):
-        # define model
-        self.l1 = nn.Linear(200, 10)
+# define what happens for training here
+def training_step(self, data_batch, batch_nb):
+    x, y = data_batch
     
-    # ---------------
-    # 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)}
+    # define your own forward and loss calculation
+    hidden_states = self.encoder(x)
      
-     # ---------------
-     # 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')
+    # 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])
         
-    # ---------------
-    # 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
+    #toy example as well
+    loss = loss / max_seq_len
+    return {'loss': loss} 
 ```
-### 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 | 
+**Or as basic as CNN image classification**      
 
-#### 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  |
+```python
+# define what happens for validation here
+def validation_step(self, data_batch, batch_nb):    
+    x, y = data_batch
     
-## Optional model hooks.   
-Add these to the model whenever you want to configure training behavior.    
+    # or as basic as a CNN classification
+    out = self.forward(x)
+    loss = my_loss(out, y)
+    return {'loss': loss} 
+```
+
+**And you also decide how to collate the output of all validation steps**    
+
+```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
+```
+   
+## Tensorboard    
+Lightning is fully integrated with tensorboard.   
+
+<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>
+
+Lightning also adds a text column with all the hyperparameters for this experiment.      
+
+<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>
+
+Simply note the path you set for the Experiment    
+``` {.python}   
+from test_tube import Experiment
+from pytorch-lightning import  Trainer
+
+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)
 
 
-### Model lifecycle hooks
-Use these hooks to customize functionality
+###### 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)
+
+
+## Demo
+```bash
+# install lightning
+pip install pytorch-lightning
+
+# clone lightning for the demo
+git clone https://github.com/williamFalcon/pytorch-lightning.git
+cd examples/new_project_templates/
+
+# run demo (on cpu)
+python trainer_gpu_cluster_template.py
+```
+
+Without changing the model AT ALL, you can run the model on a single gpu, over multiple gpus, or over multiple nodes.
+```bash
+# run a grid search on two gpus
+python fully_featured_trainer.py --gpus "0;1"
+
+# run single model on multiple gpus
+python fully_featured_trainer.py --gpus "0;1" --interactive
+```
+
 
-| 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)
+

examples/new_project_templates/__init__.py

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

examples/new_project_templates/lightning_module_template.py

@@ -0,0 +1,249 @@
+import os
+from collections import OrderedDict
+import torch.nn as nn
+from torchvision.datasets import MNIST
+import torchvision.transforms as transforms
+import torch
+import torch.nn.functional as F
+from test_tube import HyperOptArgumentParser
+from torch import optim
+from torch.utils.data import DataLoader
+from torch.utils.data.distributed import DistributedSampler
+
+from pytorch_lightning.root_module.root_module import LightningModule
+
+
+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(LightningTemplateModel, self).__init__(hparams)
+
+        self.batch_size = hparams.batch_size
+
+        # build model
+        self.__build_model()
+
+    # ---------------------
+    # MODEL SETUP
+    # ---------------------
+    def __build_model(self):
+        """
+        Layout model
+        :return:
+        """
+        self.c_d1 = nn.Linear(in_features=self.hparams.in_features, out_features=self.hparams.hidden_dim)
+        self.c_d1_bn = nn.BatchNorm1d(self.hparams.hidden_dim)
+        self.c_d1_drop = nn.Dropout(self.hparams.drop_prob)
+
+        self.c_d2 = nn.Linear(in_features=self.hparams.hidden_dim, out_features=self.hparams.out_features)
+
+    # ---------------------
+    # TRAINING
+    # ---------------------
+    def forward(self, x):
+        """
+        No special modification required for lightning, define as you normally would
+        :param x:
+        :return:
+        """
+
+        x = self.c_d1(x)
+        x = torch.tanh(x)
+        x = self.c_d1_bn(x)
+        x = self.c_d1_drop(x)
+
+        x = self.c_d2(x)
+        logits = F.log_softmax(x, dim=1)
+
+        return logits
+
+    def loss(self, labels, logits):
+        nll = F.nll_loss(logits, labels)
+        return nll
+
+    def training_step(self, data_batch, batch_i):
+        """
+        Lightning calls this inside the training loop
+        :param data_batch:
+        :return:
+        """
+        # forward pass
+        x, y = data_batch
+        x = x.view(x.size(0), -1)
+        y_hat = self.forward(x)
+
+        # calculate loss
+        loss_val = self.loss(y, y_hat)
+
+        output = OrderedDict({
+            'loss': loss_val
+        })
+
+        # can also return just a scalar instead of a dict (return loss_val)
+        return output
+
+    def validation_step(self, data_batch, batch_i):
+        """
+        Lightning calls this inside the validation loop
+        :param data_batch:
+        :return:
+        """
+        x, y = data_batch
+        x = x.view(x.size(0), -1)
+        y_hat = self.forward(x)
+
+        loss_val = self.loss(y, y_hat)
+
+        # acc
+        labels_hat = torch.argmax(y_hat, dim=1)
+        val_acc = torch.sum(y == labels_hat).item() / (len(y) * 1.0)
+
+        output = OrderedDict({
+            'val_loss': loss_val,
+            'val_acc': torch.tensor(val_acc).cuda(loss_val.device.index),
+        })
+
+        # can also return just a scalar instead of a dict (return loss_val)
+        return output
+
+    def validation_end(self, outputs):
+        """
+        Called at the end of validation to aggregate outputs
+        :param outputs: list of individual outputs of each validation step
+        :return:
+        """
+        # if returned a scalar from validation_step, outputs is a list of tensor scalars
+        # we return just the average in this case (if we want)
+        # return torch.stack(outputs).mean()
+
+        val_loss_mean = 0
+        val_acc_mean = 0
+        for output in outputs:
+            val_loss_mean += output['val_loss']
+            val_acc_mean += output['val_acc']
+
+        val_loss_mean /= len(outputs)
+        val_acc_mean /= len(outputs)
+        tqdm_dic = {'val_loss': val_loss_mean.item(), 'val_acc': val_acc_mean.item()}
+        return tqdm_dic
+
+    def update_tng_log_metrics(self, logs):
+        return logs
+
+    # ---------------------
+    # MODEL SAVING
+    # ---------------------
+    def get_save_dict(self):
+        checkpoint = {'state_dict': self.state_dict()}
+        return checkpoint
+
+    def load_model_specific(self, checkpoint):
+        self.load_state_dict(checkpoint['state_dict'])
+        pass
+
+    # ---------------------
+    # TRAINING SETUP
+    # ---------------------
+    def configure_optimizers(self):
+        """
+        return whatever optimizers we want here
+        :return: list of optimizers
+        """
+        optimizer = optim.Adam(self.parameters(), lr=self.hparams.learning_rate)
+        return [optimizer]
+
+    def __dataloader(self, train):
+        # init data generators
+        transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,), (1.0,))])
+        dataset = MNIST(root=self.hparams.data_root, train=train, transform=transform, download=True)
+
+        # when using multi-node we need to add the datasampler
+        train_sampler = None
+        batch_size = self.hparams.batch_size
+
+        try:
+            if self.on_gpu:
+                train_sampler = DistributedSampler(dataset, rank=self.trainer.proc_rank)
+                batch_size = batch_size // self.trainer.world_size  # scale batch size
+        except Exception as e:
+            pass
+
+        should_shuffle = train_sampler is None
+        loader = DataLoader(
+            dataset=dataset,
+            batch_size=batch_size,
+            shuffle=should_shuffle,
+            sampler=train_sampler
+        )
+
+        return loader
+
+    @property
+    def tng_dataloader(self):
+        if self._tng_dataloader is None:
+            try:
+                self._tng_dataloader = self.__dataloader(train=True)
+            except Exception as e:
+                print(e)
+                raise e
+        return self._tng_dataloader
+
+    @property
+    def val_dataloader(self):
+        if self._val_dataloader is None:
+            try:
+                self._val_dataloader = self.__dataloader(train=False)
+            except Exception as e:
+                print(e)
+                raise e
+        return self._val_dataloader
+
+    @property
+    def test_dataloader(self):
+        if self._test_dataloader is None:
+            try:
+                self._test_dataloader = self.__dataloader(train=False)
+            except Exception as e:
+                print(e)
+                raise e
+        return self._test_dataloader
+
+    @staticmethod
+    def add_model_specific_args(parent_parser, root_dir):
+        """
+        Parameters you define here will be available to your model through self.hparams
+        :param parent_parser:
+        :param root_dir:
+        :return:
+        """
+        parser = HyperOptArgumentParser(strategy=parent_parser.strategy, parents=[parent_parser])
+
+        # param overwrites
+        # parser.set_defaults(gradient_clip=5.0)
+
+        # network params
+        parser.opt_list('--drop_prob', default=0.2, options=[0.2, 0.5], type=float, tunable=False)
+        parser.add_argument('--in_features', default=28*28, type=int)
+        parser.add_argument('--out_features', default=10, type=int)
+        parser.add_argument('--hidden_dim', default=50000, type=int) # use 500 for CPU, 50000 for GPU to see speed difference
+
+        # data
+        parser.add_argument('--data_root', default=os.path.join(root_dir, 'mnist'), type=str)
+
+        # training params (opt)
+        parser.opt_list('--learning_rate', default=0.001*8, type=float, options=[0.0001, 0.0005, 0.001, 0.005],
+                        tunable=False)
+        parser.opt_list('--optimizer_name', default='adam', type=str, options=['adam'], tunable=False)
+
+        # if using 2 nodes with 4 gpus each the batch size here (256) will be 256 / (2*8) = 16 per gpu
+        parser.opt_list('--batch_size', default=256*8, type=int, options=[32, 64, 128, 256], tunable=False,
+                        help='batch size will be divided over all the gpus being used across all nodes')
+        return parser

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_dp_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/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/single_gpu_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='0', help='how many gpus to use in the node. -1 uses all the gpus on the node')
+    parent_parser.add_argument('--test_tube_save_path', type=str, default=test_tube_dir, help='where to save logs')
+    parent_parser.add_argument('--model_save_path', type=str, default=checkpoint_dir, help='where to save model')
+    parent_parser.add_argument('--experiment_name', type=str, default='pt_lightning_exp_a', help='test tube exp name')
+
+    # allow model to overwrite or extend args
+    parser = LightningTemplateModel.add_model_specific_args(parent_parser, root_dir)
+    hyperparams = parser.parse_args()
+
+    # ---------------------
+    # RUN TRAINING
+    # ---------------------
+    # run on HPC cluster
+    print(f'RUNNING INTERACTIVE MODE ON GPUS. gpu ids: {hyperparams.gpus}')
+    main(hyperparams)

examples/new_project_templates/trainer_cpu_template.py

@@ -0,0 +1,73 @@
+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.callbacks.pt_callbacks import EarlyStopping, ModelCheckpoint
+from docs.source.examples.example_model import ExampleModel
+
+
+def main(hparams):
+    """
+    Main training routine specific for this project
+    :param hparams:
+    :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
+    )
+
+    exp.argparse(hparams)
+    exp.save()
+
+    # build model
+    model = ExampleModel(hparams)
+
+    # callbacks
+    early_stop = EarlyStopping(
+        monitor='val_acc',
+        patience=3,
+        mode='min',
+        verbose=True,
+    )
+
+    model_save_path = '{}/{}/{}'.format(hparams.model_save_path, exp.name, exp.version)
+    checkpoint = ModelCheckpoint(
+        filepath=model_save_path,
+        save_best_only=True,
+        verbose=True,
+        monitor='val_acc',
+        mode='min'
+    )
+
+    # configure trainer
+    trainer = Trainer(
+        experiment=exp,
+        checkpoint_callback=checkpoint,
+        early_stop_callback=early_stop,
+    )
+
+    # train model
+    trainer.fit(model)
+
+
+if __name__ == '__main__':
+
+    # use default args given by lightning
+    root_dir = os.path.split(os.path.dirname(sys.modules['__main__'].__file__))[0]
+    parent_parser = HyperOptArgumentParser(strategy='random_search', add_help=False)
+    add_default_args(parent_parser, root_dir)
+
+    # allow model to overwrite or extend args
+    parser = ExampleModel.add_model_specific_args(parent_parser)
+    hyperparams = parser.parse_args()
+
+    # train model
+    main(hyperparams)

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']

pyproject.toml

@@ -0,0 +1,5 @@
+[build-system]
+requires = [
+    "setuptools",
+    "wheel",
+]

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/model_examples/bilstm.py

@@ -1,167 +0,0 @@
-import torch.nn as nn
-import numpy as np
-
-from test_tube import HyperOptArgumentParser
-import torch
-from torch.autograd import Variable
-from sklearn.metrics import confusion_matrix, f1_score
-from torch.nn import functional as F
-
-
-class BiLSTMPack(nn.Module):
-    """
-    Sample model to show how to define a template
-    """
-    def __init__(self, hparams):
-        # init superclass
-        super(BiLSTMPack, self).__init__(hparams)
-
-        self.hidden = None
-
-        # trigger tag building
-        self.ner_tagset = {'O': 0, 'I-Bio': 1}
-        self.nb_tags = len(self.ner_tagset)
-
-        # build model
-        print('building model...')
-        if hparams.model_load_weights_path is None:
-            self.__build_model()
-            print('model built')
-        else:
-            self = BiLSTMPack.load(hparams.model_load_weights_path, hparams.on_gpu, hparams)
-            print('model loaded from: {}'.format(hparams.model_load_weights_path))
-
-    def __build_model(self):
-        """
-        Layout model
-        :return:
-        """
-        # design the number of final units
-        self.output_dim = self.hparams.nb_lstm_units
-
-        # when it's bidirectional our weights double
-        if self.hparams.bidirectional:
-            self.output_dim *= 2
-
-        # total number of words
-        total_words = len(self.tng_dataloader.dataset.words_token_to_idx)
-
-        # word embeddings
-        self.word_embedding = nn.Embedding(
-            num_embeddings=total_words + 1,
-            embedding_dim=self.hparams.embedding_dim,
-            padding_idx=0
-        )
-
-        # design the LSTM
-        self.lstm = nn.LSTM(
-            self.hparams.embedding_dim,
-            self.hparams.nb_lstm_units,
-            num_layers=self.hparams.nb_lstm_layers,
-            bidirectional=self.hparams.bidirectional,
-            dropout=self.hparams.drop_prob,
-            batch_first=True,
-        )
-
-        # map to tag space
-        self.fc_out = nn.Linear(self.output_dim, self.out_dim)
-        self.hidden_to_tag = nn.Linear(self.output_dim, self.nb_tags)
-
-
-    def init_hidden(self, batch_size):
-
-        # the weights are of the form (nb_layers * 2 if bidirectional, batch_size, nb_lstm_units)
-        mult = 2 if self.hparams.bidirectional else 1
-        hidden_a = torch.randn(self.hparams.nb_layers * mult, batch_size, self.nb_rnn_units)
-        hidden_b = torch.randn(self.hparams.nb_layers * mult, batch_size, self.nb_rnn_units)
-
-        if self.hparams.on_gpu:
-            hidden_a = hidden_a.cuda()
-            hidden_b = hidden_b.cuda()
-
-        hidden_a = Variable(hidden_a)
-        hidden_b = Variable(hidden_b)
-
-        return (hidden_a, hidden_b)
-
-    def forward(self, model_in):
-        # layout data (expand it, etc...)
-        # x = sequences
-        x, seq_lengths = model_in
-        batch_size, seq_len = x.size()
-
-        # reset RNN hidden state
-        self.hidden = self.init_hidden(batch_size)
-
-        # embed
-        x = self.word_embedding(x)
-
-        # run through rnn using packed sequences
-        x = torch.nn.utils.rnn.pack_padded_sequence(x, seq_lengths, batch_first=True)
-        x, self.hidden = self.lstm(x, self.hidden)
-        x, _ = torch.nn.utils.rnn.pad_packed_sequence(x, batch_first=True)
-
-        # if asked for only last state, use the h_n which is the same as out(t=n)
-        if not self.return_sequence:
-            # pull out hidden states
-            # h_n = (nb_directions * nb_layers, batch_size, emb_size)
-            nb_directions = 2 if self.bidirectional else 1
-            (h_n, _) = self.hidden
-
-            # reshape to make indexing easier
-            # forward = 0, backward = 1 (of nb_directions)
-            h_n = h_n.view(self.nb_layers, nb_directions, batch_size, self.nb_rnn_units)
-
-            # pull out last forward
-            forward_h_n = h_n[-1, 0, :, :]
-            x = forward_h_n
-
-            # if bidirectional, also pull out the last hidden of backward network
-            if self.bidirectional:
-                backward_h_n = h_n[-1, 1, :, :]
-                x = torch.cat([forward_h_n, backward_h_n], dim=1)
-
-        # project to tag space
-        x = x.contiguous()
-        x = x.view(-1, self.output_dim)
-        x = self.hidden_to_tag(x)
-
-        return x
-
-    def loss(self, model_out):
-        # cross entropy loss
-        logits, y = model_out
-        y, y_lens = y
-
-        # flatten y and logits
-        y = y.view(-1)
-        logits = logits.view(-1, self.nb_tags)
-
-        # calculate a mask to remove padding tokens
-        mask = (y >= 0).float()
-
-        # count how many tokens we have
-        num_tokens = int(torch.sum(mask).data[0])
-
-        # pick the correct values and mask out
-        logits = logits[range(logits.shape[0]), y] * mask
-
-        # compute the ce loss
-        ce_loss = -torch.sum(logits)/num_tokens
-
-        return ce_loss
-
-    def pull_out_last_embedding(self, x, seq_lengths, batch_size, on_gpu):
-        # grab only the last activations from the non-padded ouput
-        x_last = torch.zeros([batch_size, 1, x.size(-1)])
-        for i, seq_len in enumerate(seq_lengths):
-            x_last[i, :, :] = x[i, seq_len-1, :]
-
-        # put on gpu when requested
-        if on_gpu:
-            x_last = x_last.cuda()
-
-        # turn into torch var
-        x_last = Variable(x_last)
-
-        return x_last

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,47 +1,138 @@
-import torch
-import tqdm
-import numpy as np
-from research_lib.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 research_lib.root_module.model_saving import TrainerIO
+import warnings
+import os
+import pdb
+
+import torch
+from torch.utils.data.distributed import DistributedSampler
 from torch.optim.lr_scheduler import MultiStepLR
+import torch.multiprocessing as mp
+import torch.distributed as dist
+import numpy as np
+import tqdm
+
+from pytorch_lightning.root_module.memory import get_gpu_memory_map
+from pytorch_lightning.root_module.model_saving import TrainerIO
+from pytorch_lightning.pt_overrides.override_data_parallel import LightningDistributedDataParallel, LightningDataParallel
+
+
+try:
+    from apex import amp
+    APEX_AVAILABLE = True
+except ModuleNotFoundError:
+    APEX_AVAILABLE = False
+
+
+def reduce_distributed_output(output, nb_gpus):
+    if nb_gpus <= 1:
+        return output
+
+    # when using DP, we get one output per gpu
+    # average outputs and return
+    if type(output) is torch.Tensor:
+        return output.mean()
+
+    for k, v in output.items():
+        # recurse on nested dics
+        if isinstance(output[k], dict):
+            output[k] = reduce_distributed_output(output[k], nb_gpus)
+
+        # reduce only metrics that have the same nb of gpus
+        elif output[k].size(0) == nb_gpus:
+            reduced = torch.mean(output[k])
+            output[k] = reduced
+    return output
 
 
 class Trainer(TrainerIO):
 
     def __init__(self,
                  experiment,
-                 cluster,
-                 checkpoint_callback, early_stop_callback,
+                 early_stop_callback=None,
+                 checkpoint_callback=None,
+                 gradient_clip=0,
+                 cluster=None,
                  process_position=0,
                  current_gpu_name=0,
-                 on_gpu=False,
-                 enable_tqdm=True,
-                 overfit_pct=None,
+                 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=False,
-                 enable_early_stop=True, max_nb_epochs=5, 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,
+                 accumulate_grad_batches=1,
+                 max_nb_epochs=1000, min_nb_epochs=1,
+                 train_percent_check=1.0, val_percent_check=1.0, test_percent_check=1.0,
+                 val_check_interval=0.95,
+                 log_save_interval=100, add_log_row_interval=10,
                  lr_scheduler_milestones=None,
+                 distributed_backend='dp',
+                 use_amp=False,
+                 print_nan_grads=False,
+                 print_weights_summary=True,
+                 amp_level='O2',
                  nb_sanity_val_steps=5):
 
+        """
+
+        :param experiment: Test-tube experiment
+        :param early_stop_callback: from pytorch_lightning import EarlyStopping
+        :param checkpoint_callback: from pytorch_lightning import Checkpoint
+        :param gradient_clip:
+        :param cluster:
+        :param process_position:
+        :param current_gpu_name:
+        :param nb_gpu_nodes:
+        :param gpus:
+        :param progress_bar:
+        :param overfit_pct:
+        :param track_grad_norm:
+        :param check_val_every_n_epoch:
+        :param fast_dev_run:
+        :param accumulate_grad_batches:
+        :param max_nb_epochs:
+        :param min_nb_epochs:
+        :param train_percent_check:
+        :param val_percent_check:
+        :param test_percent_check:
+        :param val_check_interval:
+        :param log_save_interval:
+        :param add_log_row_interval:
+        :param lr_scheduler_milestones:
+        :param distributed_backend: 'np' to use DistributedParallel, 'ddp' to use DistributedDataParallel
+        :param use_amp:
+        :param print_nan_grads:
+        :param print_weights_summary:
+        :param amp_level:
+        :param nb_sanity_val_steps:
+        """
+
         # Transfer params
+        self.nb_gpu_nodes = nb_gpu_nodes
+        self.gradient_clip = gradient_clip
         self.check_val_every_n_epoch = check_val_every_n_epoch
-        self.enable_early_stop = enable_early_stop
+        self.enable_early_stop = early_stop_callback is not None
         self.track_grad_norm = track_grad_norm
         self.fast_dev_run = fast_dev_run
-        self.on_gpu = 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
         self.process_position = process_position
         self.current_gpu_name = current_gpu_name
+        self.print_weights_summary = print_weights_summary
         self.checkpoint_callback = checkpoint_callback
-        self.checkpoint_callback.save_function = self.save_checkpoint
+
+        if self.checkpoint_callback is not None:
+            self.checkpoint_callback.save_function = self.save_checkpoint
+
         self.early_stop = early_stop_callback
         self.model = None
         self.max_nb_epochs = max_nb_epochs
@@ -51,6 +142,36 @@ 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
+        self.use_ddp = False
+        self.use_dp = False
+
+        # 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])
+            print(f'VISIBLE GPUS: {os.environ["CUDA_VISIBLE_DEVICES"]}')
+
+        # make DP and DDP mutually exclusive
+        # single GPU will also use DP with devices=[0]
+        have_gpus = self.data_parallel_device_ids is not None and len(self.data_parallel_device_ids) > 0
+        if have_gpus:
+            self.use_dp = distributed_backend == 'dp'
+            self.use_ddp = distributed_backend == 'ddp'
+
+        # process info
+        self.proc_rank = 0
 
         # training state
         self.optimizers = None
@@ -73,6 +194,25 @@ 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))
 
+        # 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)
+
+    @property
+    def data_parallel(self):
+        return self.use_dp or self.use_ddp
+
     def __determine_data_use_amount(self, train_percent_check, val_percent_check, test_percent_check, overfit_pct):
         """
         Use less data for debugging purposes
@@ -85,20 +225,27 @@ class Trainer(TrainerIO):
             self.val_percent_check = overfit_pct
             self.test_percent_check = overfit_pct
 
+    def __get_model(self):
+        return self.model.module if self.data_parallel else self.model
+
     def __is_function_implemented(self, f_name):
-        f_op = getattr(self, f_name, None)
+        model = self.__get_model()
+        f_op = getattr(model, f_name, None)
         return callable(f_op)
 
     @property
     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):
@@ -110,24 +257,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):
@@ -138,8 +288,6 @@ class Trainer(TrainerIO):
         :param max_batches: Scalar
         :return:
         """
-        print('validating...')
-
         # enable eval mode
         model.zero_grad()
         model.eval()
@@ -151,33 +299,45 @@ 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.use_ddp:
+                output = model(data_batch, batch_i)
+            elif self.use_dp:
+                output = model(data_batch, batch_i)
+                output = reduce_distributed_output(output, len(self.data_parallel_device_ids))
+
+            else:
+                output = model.validation_step(data_batch, batch_i)
+
             outputs.append(output)
 
             # batch done
-            if self.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):
@@ -190,22 +350,188 @@ class Trainer(TrainerIO):
         self.test_dataloader = model.test_dataloader
         self.val_dataloader = model.val_dataloader
 
+        if self.use_ddp and not isinstance(self.tng_dataloader.sampler, DistributedSampler):
+            msg = '''
+            when using multiple gpus and multiple nodes you must pass a DistributedSampler to DataLoader(sampler).
+            
+            ie: this:
+            dataset = myDataset()
+            dataloader = Dataloader(dataset)
+            
+            becomes:
+            dataset = myDataset()
+            dist_sampler = torch.utils.data.distributed.DistributedSampler(dataset)
+            dataloader = Dataloader(dataset, sampler=dist_sampler)
+            '''
+            raise Exception(msg)
+
     # -----------------------------
     # MODEL TRAINING
     # -----------------------------
     def fit(self, model):
-        self.model = model
 
-        # transfer data loaders from model
-        self.__get_dataloaders(model)
+        # when using multi-node or DDP within a node start each module in a separate process
+        if self.use_ddp:
+            # must copy only the meta of the exp so it survives pickle/unpickle when going to new process
+            self.experiment = self.experiment.get_meta_copy()
 
-        # init training constants
-        self.__layout_bookeeping()
+            # whenever we have the correct number of tasks, we let slurm manage processes
+            # otherwise we launch the required number of processes
+            nb_slurm_tasks = int(os.environ['SLURM_NTASKS'])
+            nb_requested_gpus = len(self.data_parallel_device_ids)
+            is_slurm_managing_tasks = nb_slurm_tasks == nb_requested_gpus
+            if is_slurm_managing_tasks:
+                task = int(os.environ['SLURM_LOCALID'])
+                self.ddp_train(task, model)
+            else:
+                msg = f"""
+                You requested {nb_requested_gpus} GPUs but launched {nb_slurm_tasks} slurm tasks. 
+                We will launch {nb_requested_gpus} processes for you. 
+                We recommend you let slurm manage the processes by setting: --ntasks-per-node={nb_requested_gpus}
+                """
+                warnings.warn(msg)
+                mp.spawn(self.ddp_train, nprocs=len(self.data_parallel_device_ids), args=(model, ))
+
+        # 1 gpu or dp option triggers training using DP module
+        # easier to avoid NCCL issues
+        elif self.use_dp:
+            self.dp_train(model)
+
+        # ON CPU
+        else:
+            # CHOOSE OPTIMIZER
+            # filter out the weights that were done on gpu so we can load on good old cpus
+            self.optimizers = model.configure_optimizers()
+
+            # 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, model):
 
         # CHOOSE OPTIMIZER
         # filter out the weights that were done on gpu so we can load on good old cpus
         self.optimizers = model.configure_optimizers()
 
+        model.cuda(self.data_parallel_device_ids[0])
+        model = LightningDataParallel(model, device_ids=self.data_parallel_device_ids)
+
+        # 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 ddp_train(self, gpu_nb, model):
+        """
+        Entry point into a DP thread
+        :param gpu_nb:
+        :param model:
+        :param cluster_obj:
+        :return:
+        """
+        # node rank using relative slurm id
+        # otherwise default to node rank 0
+        try:
+            node_rank = int(os.environ['SLURM_NODEID'])
+        except KeyError as e:
+            node_rank = 0
+
+        # recover original exp before went into process
+        # init in write mode only on proc 0
+        self.experiment.debug = self.proc_rank > 0
+        self.experiment = self.experiment.get_non_ddp_exp()
+
+        # show progbar only on prog_rank 0
+        self.prog_bar = self.prog_bar and 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
+        # where to store ip_table
+        self.__init_tcp_connection()
+
+        # CHOOSE OPTIMIZER
+        # filter out the weights that were done on gpu so we can load on good old cpus
+        self.optimizers = model.configure_optimizers()
+
+        # MODEL
+        # copy model to each gpu
+        torch.cuda.set_device(gpu_nb)
+        model.cuda(gpu_nb)
+
+        # AMP
+        # run through amp wrapper before going to distributed DP
+        if self.use_amp:
+            # An example
+            model, optimizers = amp.initialize(
+                model, self.optimizers, opt_level=self.amp_level,
+            )
+            self.optimizers = optimizers
+
+        model = LightningDistributedDataParallel(model, device_ids=[gpu_nb])
+
+        # continue training routine
+        self.__run_pretrain_routine(model)
+
+    def __init_tcp_connection(self):
+        """
+        Connect all procs in the world using the env:// init
+        Use the first node as the root address
+        :param port:
+        :param tries:
+        :return:
+        """
+        try:
+            port = os.environ['MASTER_PORT']
+        except Exception as e:
+            port = 12910
+            os.environ['MASTER_PORT'] = f'{port}'
+
+        try:
+            root_node = os.environ['SLURM_NODELIST'].split(' ')[0]
+        except Exception as e:
+            root_node = '127.0.0.2'
+
+        os.environ['MASTER_ADDR'] = root_node
+
+        sleep(self.proc_rank*0.5)
+        dist.init_process_group("nccl", rank=self.proc_rank, world_size=self.world_size)
+
+    def __run_pretrain_routine(self, model):
+        """
+        Sanity check a few things before starting actual training
+        :param model:
+        :return:
+        """
+        ref_model = model
+        if self.data_parallel:
+            ref_model = model.module
+
+        ref_model.trainer = self
+
+        # set local properties on the model
+        ref_model.on_gpu = self.on_gpu
+
+        # transfer data loaders from model
+        self.__get_dataloaders(ref_model)
+
+        # init training constants
+        self.__layout_bookeeping()
+
         # add lr schedulers
         if self.lr_scheduler_milestones is not None:
             for optimizer in self.optimizers:
@@ -213,24 +539,28 @@ class Trainer(TrainerIO):
                 self.lr_schedulers.append(scheduler)
 
         # print model summary
-        model.summarize()
+        if self.proc_rank == 0 and self.print_weights_summary:
+            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
-        self.enable_auto_hpc_walltime_manager()
+        if self.cluster is not None:
+            self.enable_auto_hpc_walltime_manager()
 
         # ---------------------------
         # CORE TRAINING LOOP
         # ---------------------------
+        self.model = model
         self.__train()
 
     def __train(self):
@@ -240,24 +570,28 @@ class Trainer(TrainerIO):
             for lr_scheduler in self.lr_schedulers:
                 lr_scheduler.step()
 
-            self.model.current_epoch = epoch_nb
+            model = self.__get_model()
+            model.current_epoch = epoch_nb
 
             # hook
             if self.__is_function_implemented('on_epoch_start'):
-                self.model.on_epoch_start()
+                model = self.__get_model()
+                model.on_epoch_start()
 
             self.current_epoch = epoch_nb
             self.total_batches = self.nb_tng_batches + self.nb_val_batches
             self.batch_loss_value = 0  # accumulated grads
 
             # init progbar when requested
-            if self.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.__get_model()
+                model.global_step = self.global_step
 
                 # stop when the flag is changed or we've gone past the amount requested in the batches
                 self.total_batch_nb += 1
@@ -268,25 +602,28 @@ class Trainer(TrainerIO):
                 # ---------------
                 # RUN TRAIN STEP
                 # ---------------
-                self.__run_tng_batch(data_batch)
+                batch_result = self.__run_tng_batch(data_batch, batch_nb)
+                early_stop_epoch = batch_result == -1
 
                 # ---------------
                 # RUN VAL STEP
                 # ---------------
                 is_val_check_batch = (batch_nb + 1) % self.val_check_batch == 0
-                if self.fast_dev_run or is_val_check_batch:
+                if self.fast_dev_run or is_val_check_batch or early_stop_epoch:
                     self.__run_validation()
 
                 # when batch should be saved
-                if (batch_nb + 1) % self.log_save_interval == 0:
-                    self.experiment.save()
+                if (batch_nb + 1) % self.log_save_interval == 0 or early_stop_epoch:
+                    if self.proc_rank == 0:
+                        self.experiment.save()
 
                 # when metrics should be logged
-                if batch_nb % self.add_log_row_interval == 0:
+                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)
+                    model = self.__get_model()
+                    metrics = model.update_tng_log_metrics(self.__tng_tqdm_dic)
 
                     # add gpu memory
                     if self.on_gpu:
@@ -295,54 +632,127 @@ class Trainer(TrainerIO):
 
                     # add norms
                     if self.track_grad_norm > 0:
-                        grad_norm_dic = self.model.grad_norm(self.track_grad_norm)
+                        model = self.__get_model()
+                        grad_norm_dic = model.grad_norm(self.track_grad_norm)
                         metrics.update(grad_norm_dic)
 
+                    if self.__is_function_implemented('on_tng_metrics'):
+                        model.on_tng_metrics(metrics)
+
                     # log metrics
-                    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.__get_model()
+                    model.on_batch_end()
+
+                # end epoch early
+                if early_stop_epoch:
+                    break
 
             # hook
             if self.__is_function_implemented('on_epoch_end'):
-                self.model.on_epoch_end()
+                model = self.__get_model()
+                model.on_epoch_end()
 
             # early stopping
-            if self.enable_early_stop:
+            met_min_epochs = epoch_nb > self.min_nb_epochs
+            if self.enable_early_stop and met_min_epochs:
                 should_stop = self.early_stop_callback.on_epoch_end(epoch=epoch_nb, logs=self.__tng_tqdm_dic)
-                met_min_epochs = epoch_nb > self.min_nb_epochs
 
                 # stop training
                 stop = should_stop and met_min_epochs
                 if stop:
                     return
 
-    def __run_tng_batch(self, data_batch):
+    def __metrics_to_scalars(self, metrics, blacklist=[]):
+        new_metrics = {}
+        for k, v in metrics.items():
+            if type(v) is torch.Tensor:
+                v = v.item()
+
+            if type(v) is dict:
+                v = self.__metrics_to_scalars(v)
+
+            if k not in blacklist:
+                new_metrics[k] = float(v)
+
+        return new_metrics
+
+    def __log_vals_blacklist(self):
+        """avoid logging some vals lightning uses to maintain state"""
+        blacklist = {'batch_nb', 'v_nb', 'gpu'}
+        return blacklist
+
+    def __run_tng_batch(self, data_batch, batch_nb):
         if data_batch is None:
-            return
+            return 0
 
         # hook
         if self.__is_function_implemented('on_batch_start'):
-            self.model.on_batch_start()
+            model_ref = self.__get_model()
+            response = model_ref.on_batch_start(data_batch)
 
-        if self.enable_tqdm:
+            if response == -1:
+                return -1
+
+        if self.progress_bar:
             self.prog_bar.update(1)
 
         # forward pass
         # return a scalar value and a dic with tqdm metrics
-        loss, model_specific_tqdm_metrics_dic = self.model.training_step(data_batch)
+        if self.use_ddp:
+            output = self.model(data_batch, batch_nb)
+        elif self.use_dp:
+            output = self.model(data_batch, batch_nb)
+            output = reduce_distributed_output(output, len(self.data_parallel_device_ids))
+        else:
+            output = self.model.training_step(data_batch, batch_nb)
+
+        try:
+            model_specific_tqdm_metrics_dic = output['tqdm_metrics']
+        except Exception as e:
+            model_specific_tqdm_metrics_dic = {}
+
+        # if output dict doesn't have the keyword loss
+        # then assume the output=loss if scalar
+        try:
+            loss = output['loss']
+        except Exception as e:
+            if type(output) is torch.Tensor:
+                loss = output
+
         self.__add_tqdm_metrics(model_specific_tqdm_metrics_dic)
 
         # backward pass
-        loss.backward()
+        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.__get_model()
+            for param in model.parameters():
+                print(param.grad.float().sum())
+
+        # avoid memory leaks
         self.batch_loss_value += loss.item()
 
         # gradient update with accumulated gradients
         if (self.batch_nb + 1) % self.accumulate_grad_batches == 0:
 
+            # clip gradients
+            if self.gradient_clip > 0:
+                model = self.__get_model()
+                torch.nn.utils.clip_grad_norm(model.parameters(), self.gradient_clip)
+
             # update gradients across all optimizers
             for optimizer in self.optimizers:
                 optimizer.step()
@@ -359,14 +769,17 @@ 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)
 
         # activate batch end hook
         if self.__is_function_implemented('on_batch_end'):
-            self.model.on_batch_end()
+            model = self.__get_model()
+            model.on_batch_end()
+
+        return 0
 
     def __run_validation(self):
         # decide if can check epochs
@@ -379,7 +792,8 @@ class Trainer(TrainerIO):
         try:
             # hook
             if self.__is_function_implemented('on_pre_performance_check'):
-                self.model.on_pre_performance_check()
+                model = self.__get_model()
+                model.on_pre_performance_check()
 
             # use full val set on end of epoch
             # use a small portion otherwise
@@ -393,17 +807,19 @@ class Trainer(TrainerIO):
 
             # hook
             if self.__is_function_implemented('on_post_performance_check'):
-                self.model.on_post_performance_check()
+                model = self.__get_model()
+                model.on_post_performance_check()
 
         except Exception as e:
             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 and self.checkpoint_callback:
+            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,188 @@
+from torch.nn import DataParallel
+from torch.nn.parallel import DistributedDataParallel
+import itertools
+from itertools import chain
+
+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 forward(self, *inputs, **kwargs):
+        if not self.device_ids:
+            return self.module(*inputs, **kwargs)
+
+        for t in chain(self.module.parameters(), self.module.buffers()):
+            if t.device != self.src_device_obj:
+                raise RuntimeError("module must have its parameters and buffers "
+                                   "on device {} (device_ids[0]) but found one of "
+                                   "them on device: {}".format(self.src_device_obj, t.device))
+
+        inputs, kwargs = self.scatter(inputs, kwargs, self.device_ids)
+        if len(self.device_ids) == 1:
+            # lightning
+            if self.module.training:
+                return self.module.training_step(*inputs[0], **kwargs[0])
+            else:
+                return self.module.validation_step(*inputs[0], **kwargs[0])
+
+        replicas = self.replicate(self.module, self.device_ids[:len(inputs)])
+        outputs = self.parallel_apply(replicas, inputs, kwargs)
+        return self.gather(outputs, self.output_device)
+
+
+    def parallel_apply(self, replicas, inputs, kwargs):
+        return parallel_apply(replicas, inputs, kwargs, self.device_ids[:len(replicas)])
+
+
+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/hooks.py

@@ -1,7 +1,7 @@
 import torch
 
 class ModelHooks(torch.nn.Module):
-    def on_batch_start(self):
+    def on_batch_start(self, data_batch):
         pass
 
     def on_batch_end(self):
@@ -18,3 +18,7 @@ class ModelHooks(torch.nn.Module):
 
     def on_post_performance_check(self):
         pass
+
+    def on_tng_metrics(self, metrics):
+        pass
+

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, LightningDataParallel
 
 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,9 @@ class TrainerIO(object):
         checkpoint['optimizer_states'] = optimizer_states
 
         # request what to save from the model
-        checkpoint_dict = self.model.get_save_dict()
+        is_dp_module = type(self.model) is LightningDistributedDataParallel or type(self.model) is LightningDataParallel
+        model = self.model.module if is_dp_module else self.model
+        checkpoint_dict = model.get_save_dict()
 
         # merge trainer and model saving items
         checkpoint.update(checkpoint_dict)
@@ -88,6 +108,7 @@ class TrainerIO(object):
         self.early_stop_callback.wait = checkpoint['early_stop_callback_wait']
         self.early_stop_callback.patience = checkpoint['early_stop_callback_patience']
         self.global_step = checkpoint['global_step']
+        self.current_epoch = checkpoint['epoch']
 
         # restore the optimizers
         optimizer_states = checkpoint['optimizer_states']
@@ -98,15 +119,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()
 
@@ -125,10 +155,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,32 +9,28 @@ 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.on_gpu = hparams.on_gpu
+
         self.dtype = torch.FloatTensor
         self.exp_save_path = None
         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
 
         # 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...')
-            self.dtype = torch.cuda.FloatTensor
-            torch.set_default_tensor_type('torch.cuda.FloatTensor')
-
     def forward(self, *args, **kwargs):
         """
         Expand model in into whatever you need.
@@ -44,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:
@@ -60,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:
@@ -82,7 +78,7 @@ class RootModule(GradInformation, ModelIO, OptimizerConfig, ModelHooks):
         :param logs:
         :return:
         """
-        raise NotImplementedError
+        return logs
 
     def loss(self, *args, **kwargs):
         """
@@ -96,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():
@@ -136,26 +129,30 @@ class RootModule(GradInformation, ModelIO, OptimizerConfig, ModelHooks):
     def get_process_position(gpus):
         try:
             current_gpu = os.environ["CUDA_VISIBLE_DEVICES"]
-            gpu_ids = gpus.split(';')
+            gpu_ids = gpus.split(',')
             process_position = gpu_ids.index(current_gpu)
             return process_position, current_gpu
         except Exception as e:
             return 0, 0
 
     @classmethod
-    def load_from_metrics(cls, weights_path, tags_csv, on_gpu):
+    def load_from_metrics(cls, weights_path, tags_csv, on_gpu, map_location=None):
         """
         Primary way of loading model from csv weights path
         :param weights_path:
         :param tags_csv:
         :param on_gpu:
+        :param map_location: dic for mapping storage {'cuda:1':'cuda:0'}
         :return:
         """
         hparams = load_hparams_from_tags_csv(tags_csv)
         hparams.__setattr__('on_gpu', on_gpu)
 
         if on_gpu:
-            checkpoint = torch.load(weights_path)
+            if map_location is not None:
+                checkpoint = torch.load(weights_path, map_location=map_location)
+            else:
+                checkpoint = torch.load(weights_path)
         else:
             checkpoint = torch.load(weights_path, map_location=lambda storage, loc: storage)
 

pytorch_lightning/trainer_main.py

@@ -8,8 +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)
@@ -98,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,4 +1,6 @@
-def add_default_args(parser, root_dir, possible_model_names, rand_seed):
+import pdb
+
+def add_default_args(parser, root_dir, rand_seed=None, possible_model_names=None):
 
     # tng, test, val check intervals
     parser.add_argument('--eval_test_set', dest='eval_test_set', action='store_true', help='true = run test set also')
@@ -32,21 +34,26 @@ def add_default_args(parser, root_dir, possible_model_names, rand_seed):
 
     # model paths
     parser.add_argument('--model_load_weights_path', default=None, type=str)
-    parser.add_argument('--model_name', default='', help=','.join(possible_model_names))
+
+    if possible_model_names is not 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')
@@ -58,10 +65,12 @@ def add_default_args(parser, root_dir, possible_model_names, rand_seed):
     parser.add_argument('--overfit', default=-1, type=float, help='% of dataset to use with this option. float, or -1 for none')
 
     # debug args
-    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')
+    if rand_seed is not None:
+        parser.add_argument('--random_seed', default=rand_seed, type=int)
+
+    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)

pytorch_lightning/utils/embeddings.py

@@ -13,11 +13,7 @@ class PretrainedEmbedding(torch.nn.Embedding):
         >>> emb = PretrainedEmbedding(embedding_path='glove.840B.300d.txt',embedding_dim=300, task_vocab={'hello': 1, 'world': 2})
         >>> data = torch.Tensor([[0, 1], [0, 2]]).long()
         >>> embedded = emb(data)
-        tensor([[[ 0.0000,  0.0000,  0.0000,  ...,  0.0000,  0.0000,  0.0000],
-         [ 0.2523,  0.1018, -0.6748,  ...,  0.1787, -0.5192,  0.3359]],
-
-        [[ 0.0000,  0.0000,  0.0000,  ...,  0.0000,  0.0000,  0.0000],
-         [-0.0067,  0.2224,  0.2771,  ...,  0.0594,  0.0014,  0.0987]]])
+        
 
 
         :param embedding_path:
@@ -37,7 +33,8 @@ class PretrainedEmbedding(torch.nn.Embedding):
         self.weight = new_emb.weight
 
         # apply freeze
-        self.weight.requires_grad = not freeze
+        should_freeze = not freeze
+        self.weight.requires_grad = should_freeze
 
     def __load_task_specific_embeddings(self, vocab_words, embedding_path, emb_dim, freeze):
         """
@@ -97,11 +94,11 @@ class PretrainedEmbedding(torch.nn.Embedding):
 
 if __name__ == '__main__':
     emb = PretrainedEmbedding(
-        embedding_path='/Users/waf/Developer/NGV/research-fermat/fermat/.vector_cache/glove.840B.300d.txt',
+        embedding_path='/Users/waf/Developer',
         embedding_dim=300,
         task_vocab={'hello': 1, 'world': 2}
     )
 
     data = torch.Tensor([[0, 1], [0, 2]]).long()
     embedded = emb(data)
-    print(embedded)
+    print(embedded)

pytorch_lightning/utils/plotting.py

@@ -1,4 +1,3 @@
-from matplotlib import pyplot as plt
 import numpy as np
 np.seterr(divide='ignore', invalid='ignore')
 
@@ -13,6 +12,7 @@ def plot_confusion_matrix(cm,
     This function prints and plots the confusion matrix.
     Normalization can be applied by setting `normalize=True`.
     """
+    from matplotlib import pyplot as plt
     if normalize:
         cm = cm.astype('float') / cm.sum(axis=1)[:, np.newaxis]
         print("Normalized confusion matrix")

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.cfg

@@ -0,0 +1,21 @@
+[tool:pytest]
+norecursedirs =
+    .git
+    dist
+    build
+python_files =
+    test_*.py
+doctest_plus = disabled
+addopts = --strict
+markers =
+    slow
+    remote_data
+    filterwarnings
+
+[pycodestyle]
+ignore = E731,W504
+max-line-length = 120
+
+[flake8]
+ignore = E731,W504,F401,F841
+max-line-length = 120

setup.py

@@ -2,12 +2,28 @@
 
 from setuptools import setup, find_packages
 
-setup(name='pytorch-lightning',
-      version='0.0.1',
-      description='Rapid research framework',
-      author='',
-      author_email='',
-      url='https://github.com/williamFalcon/pytorch-lightning',
-      install_requires=[],
-      packages=find_packages()
-      )
+# https://packaging.python.org/guides/single-sourcing-package-version/
+
+# http://blog.ionelmc.ro/2014/05/25/python-packaging/
+setup(
+    name="pytorch-lightning",
+    version='0.2.5.1',
+    description="The Keras for ML researchers using PyTorch",
+    author="William Falcon",
+    author_email="waf2107@columbia.edu",
+    url="https://github.com/williamFalcon/pytorch-lightning",
+    download_url="https://github.com/williamFalcon/pytorch-lightning",
+    license="MIT",
+    keywords=["deep learning", "pytorch", "AI"],
+    python_requires=">=3.5",
+    install_requires=[
+        "torch>=1.1.0",
+        "tqdm",
+        "test-tube>=0.6.7.1",
+    ],
+    packages=find_packages(),
+    long_description=open("README.md", encoding="utf-8").read(),
+    long_description_content_type='text/markdown',
+    include_package_data=True,
+    zip_safe=False,
+)

tests/README.md

@@ -1,65 +0,0 @@
-# Testing setup
-
-## A. Enable CircleCI for your project
-1. Integrate CircleCI by clicking "Set up Project" at [this link](https://circleci.com/add-projects/gh/NextGenVest).
-
-## B. Add your own tests
-1. In the /tests, emulate exactly the folder structure for your module found under /bot_seed
-2. To create a test for file ```/bot_seed/folder/example.py```:
-    - create the file ```/tests/folder/example_test.py```
-    - notice the **_test**
-    - notice the mirror path under **/tests**
-
-3. Your ```example_test.py``` file should have these main components
-
-```python
-# example.py
-
-def function_i_want_to_test(x):
-    return x*2
-
-def square(x):
-    return x*x
-
-```
-
-```python
-# example_test.py
-
-import pytest
-
-# do whatever imports you need
-from app.bot_seed.folder.example import function_i_want_to_test, square
-
-def test_function_i_want_to_test():
-    answer = function_i_want_to_test(4)
-    assert answer == 8
-
-# -----------------------------------
-# Your function must start with test_
-# -----------------------------------
-def test_square():
-    answer = square(3)
-    assert answer == 9
-
-# -----------------------------------
-# boilerplate (link this file to pytest)
-# -----------------------------------
-if __name__ == '__main__':
-    pytest.main([__file__])
-```
-
-## C. Add build passing badge
-1. Create a CircleCI status token:
-    - Go here: https://circleci.com/gh/NextGenVest/your-project-name/edit#api
-    - Click create token
-    - Select status
-    - Type "badge status"
-
-2. Get a copy of the markdown code:
-    - Go here: https://circleci.com/gh/NextGenVest/your-project-name/edit#badges
-    - Select master
-    - Select "badge status" token
-    - Select image URL
-    - Copy the image url link and change the html at the top of the root README.md file for your project
-    

tests/research_proj/sample_model_template/model_template_test.py

@@ -1,13 +0,0 @@
-import pytest
-
-"""
-Example test to show how to add a test for anything in the project.  
-Look at the README for more instructions   
-"""
-
-
-def test_cube():
-    assert 27 == 27
-
-if __name__ == '__main__':
-    pytest.main([__file__])