release v0.3.5

removed dep

.gitignore

@@ -8,6 +8,8 @@ datasets/
 model_weights/
 app/models/
 pip-wheel-metadata/
+test_tube_exp/
+tests/tests_tt_dir/
 
 # Byte-compiled / optimized / DLL files
 __pycache__/
@@ -118,4 +120,4 @@ ENV/
 .mypy_cache/
 
 # data
-mnist/
+mnist/

.readthedocs.yml

@@ -0,0 +1,19 @@
+# .readthedocs.yml
+# Read the Docs configuration file
+# See https://docs.readthedocs.io/en/stable/config-file/v2.html for details
+
+# Required
+version: 2
+
+# Build documentation with MkDocs
+mkdocs:
+  configuration: mkdocs.yml
+
+# Optionally build your docs in additional formats such as PDF and ePub
+formats: all
+
+# Optionally set the version of Python and requirements required to build your docs
+python:
+  version: 3.7
+  install:
+    - requirements: docs/doc_requirements.txt

.travis.yml

@@ -0,0 +1,16 @@
+language: python
+python:
+  - "3.7"
+# command to install dependencies
+cache: pip
+install:
+  - pip install -e .
+  - pip install -r requirements.txt
+  - pip install -U numpy
+
+# keep build from timing out
+dist: xenial
+
+# command to run tests
+script:
+  - py.test # or py.test for Python versions 3.5 and below

README.md

@@ -11,12 +11,14 @@
 </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/tree/master/tests"><img src="https://github.com/williamFalcon/pytorch-lightning/blob/master/coverage.svg"></a>
+  <a href="https://travis-ci.org/williamFalcon/pytorch-lightning"><img src="https://travis-ci.org/williamFalcon/pytorch-lightning.svg?branch=master"></a>
+  <a href="https://williamfalcon.github.io/pytorch-lightning/"><img src="https://readthedocs.org/projects/pytorch-lightning/badge/?version=latest"></a>
   <a href="https://github.com/williamFalcon/pytorch-lightning/blob/master/COPYING"><img src="https://img.shields.io/badge/License-MIT-yellow.svg"></a>
 </p>   
 
 ```bash
-pip install pytorch-lightning    
+pip install pytorch-lightning  
 ```
 
 ## Docs   
@@ -36,9 +38,26 @@ To use lightning do 2 things:
 2. [Define a LightningModel](https://github.com/williamFalcon/pytorch-lightning/blob/master/examples/new_project_templates/lightning_module_template.py).     
 
 ## What does lightning control for me?
-Everything! Except the following three things:
+Everything!    
+Except for these 6 core functions which you define:    
 
-**What happens in the training loop**    
+```{.python}
+# what to do in the training loop
+def training_step(self, data_batch, batch_nb):
+
+# what to do in the validation loop
+def validation_step(self, data_batch, batch_nb):
+
+# how to aggregate validation_step outputs
+def validation_end(self, outputs):
+
+# and your dataloaders
+def tng_dataloader():
+def val_dataloader():
+def test_dataloader():
+```
+
+**Could be as complex as seq-2-seq + attention**    
 
 ```python
 # define what happens for training here
@@ -46,25 +65,39 @@ def training_step(self, data_batch, batch_nb):
     x, y = data_batch
     
     # define your own forward and loss calculation
-    out = self.forward(x)
-    loss = my_loss(out, y)
+    hidden_states = self.encoder(x)
+     
+    # even as complex as a seq-2seq + attn model
+    # (this is just a toy, non-working example to illustrate)
+    start_token = '<SOS>'
+    last_hidden = torch.zeros(...)
+    loss = 0
+    for step in range(max_seq_len):
+        attn_context = self.attention_nn(hidden_states, start_token)
+        pred = self.decoder(start_token, attn_context, last_hidden) 
+        last_hidden = pred
+        pred = self.predict_nn(pred)
+        loss += self.loss(last_hidden, y[step])
+        
+    #toy example as well
+    loss = loss / max_seq_len
     return {'loss': loss} 
 ```
 
-**What happens in the validation loop**      
+**Or as basic as CNN image classification**      
 
 ```python
 # define what happens for validation here
 def validation_step(self, data_batch, batch_nb):    
     x, y = data_batch
     
-    # define your own forward and loss calculation
+    # or as basic as a CNN classification
     out = self.forward(x)
     loss = my_loss(out, y)
     return {'loss': loss} 
 ```
 
-**And what to do with the output of all validation batches**    
+**And you also decide how to collate the output of all validation steps**    
 
 ```python
 def validation_end(self, outputs):
@@ -84,22 +117,52 @@ def validation_end(self, 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.   
 
-## Lightning gives you options to control the following:
+<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>
 
-**Checkpointing**    
+Lightning also adds a text column with all the hyperparameters for this experiment.      
 
-- Model saving
-- Model loading 
+<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>
 
-**Computing cluster (SLURM)**    
+Simply note the path you set for the Experiment    
+``` {.python}   
+from test_tube import Experiment
+from pytorch-lightning import  Trainer
 
-- Automatic checkpointing   
-- Automatic saving, loading  
-- Running grid search on a cluster 
-- Walltime auto-resubmit   
+exp = Experiment(save_dir='/some/path')
+trainer = Trainer(experiment=exp)
+...
+```   
 
-**Debugging**  
+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)
@@ -109,7 +172,7 @@ def validation_end(self, outputs):
 - [Pring which gradients are nan](https://williamfalcon.github.io/pytorch-lightning/Trainer/debugging/#print-which-gradients-are-nan)
 
 
-**Distributed training**    
+###### Distributed training    
 
 - [16-bit mixed precision](https://williamfalcon.github.io/pytorch-lightning/Trainer/Distributed%20training/#16-bit-mixed-precision)
 - [Multi-GPU](https://williamfalcon.github.io/pytorch-lightning/Trainer/Distributed%20training/#Multi-GPU)
@@ -118,26 +181,27 @@ def validation_end(self, outputs):
 - [Self-balancing architecture](https://williamfalcon.github.io/pytorch-lightning/Trainer/Distributed%20training/#self-balancing-architecture)
 
 
-**Experiment Logging**   
+###### Experiment Logging   
 
 - [Display metrics in progress bar](https://williamfalcon.github.io/pytorch-lightning/Trainer/Logging/#display-metrics-in-progress-bar)
 - Log arbitrary metrics
 - [Log metric row every k batches](https://williamfalcon.github.io/pytorch-lightning/Trainer/Logging/#log-metric-row-every-k-batches)
-- [Process position](Logging/#process-position)
+- [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**    
+###### 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**    
+###### 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)

coverage.svg

@@ -0,0 +1,21 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<svg xmlns="http://www.w3.org/2000/svg" width="99" height="20">
+    <linearGradient id="b" x2="0" y2="100%">
+        <stop offset="0" stop-color="#bbb" stop-opacity=".1"/>
+        <stop offset="1" stop-opacity=".1"/>
+    </linearGradient>
+    <mask id="a">
+        <rect width="99" height="20" rx="3" fill="#fff"/>
+    </mask>
+    <g mask="url(#a)">
+        <path fill="#555" d="M0 0h63v20H0z"/>
+        <path fill="#4c1" d="M63 0h36v20H63z"/>
+        <path fill="url(#b)" d="M0 0h99v20H0z"/>
+    </g>
+    <g fill="#fff" text-anchor="middle" font-family="DejaVu Sans,Verdana,Geneva,sans-serif" font-size="11">
+        <text x="31.5" y="15" fill="#010101" fill-opacity=".3">coverage</text>
+        <text x="31.5" y="14">coverage</text>
+        <text x="80" y="15" fill="#010101" fill-opacity=".3">99%</text>
+        <text x="80" y="14">99%</text>
+    </g>
+</svg>

docs/LightningModule/RequiredTrainerInterface.md

@@ -1,30 +1,30 @@
-# Lightning module  
+# 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. 
+The easiest thing to do is copy [this template](../../pytorch_lightning/examples/new_project_templates/lightning_module_template.py) and modify accordingly. 
 
 Otherwise, to Define a Lightning Module, implement the following methods:
 
 **Required**:  
 
-- [training_step](LightningModule.md#training_step)   
-- [validation_step](LightningModule.md#validation_step)
-- [validation_end](LightningModule.md#validation_end)
+- [training_step](RequiredTrainerInterface.md#training_step)   
+- [validation_step](RequiredTrainerInterface.md#validation_step)
+- [validation_end](RequiredTrainerInterface.md#validation_end)
     
-- [configure_optimizers](LightningModule.md#configure_optimizers)
-- [get_save_dict](LightningModule.md#get_save_dict)
-- [load_model_specific](LightningModule.md#load_model_specific)
+- [configure_optimizers](RequiredTrainerInterface.md#configure_optimizers)
+- [get_save_dict](RequiredTrainerInterface.md#get_save_dict)
+- [load_model_specific](RequiredTrainerInterface.md#load_model_specific)
 
-- [tng_dataloader](LightningModule.md#tng_dataloader)
-- [tng_dataloader](LightningModule.md#tng_dataloader)
-- [test_dataloader](LightningModule.md#test_dataloader)
+- [tng_dataloader](RequiredTrainerInterface.md#tng_dataloader)
+- [tng_dataloader](RequiredTrainerInterface.md#tng_dataloader)
+- [test_dataloader](RequiredTrainerInterface.md#test_dataloader)
 
 **Optional**:   
 
-- [update_tng_log_metrics](LightningModule.md#update_tng_log_metrics)
-- [add_model_specific_args](LightningModule.md#add_model_specific_args)
+- [update_tng_log_metrics](RequiredTrainerInterface.md#update_tng_log_metrics)
+- [add_model_specific_args](RequiredTrainerInterface.md#add_model_specific_args)
 
 ---
 
@@ -361,7 +361,7 @@ 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 argument parser is available anywhere in your model by calling self.hparams
+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
@@ -391,4 +391,4 @@ def add_model_specific_args(parent_parser, root_dir):
     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

@@ -3,6 +3,26 @@ Lightning makes multi-gpu training and 16 bit training trivial.
 *Note:*   
 None of the flags below require changing anything about your lightningModel definition. 
 
+---
+#### Choosing a backend  
+Lightning supports two backends. DataParallel and DistributedDataParallel. Both can be used for single-node multi-GPU training.
+For multi-node training you must use DistributedDataParallel.   
+
+You can toggle between each mode by setting this flag.
+``` {.python}
+# DEFAULT uses DataParallel
+trainer = Trainer(distributed_backend='dp')
+
+# change to distributed data parallel
+trainer = Trainer(distributed_backend='ddp')
+```
+
+If you request multiple nodes, the back-end will auto-switch to ddp.
+We recommend you use DistributedDataparallel even for single-node multi-GPU training. It is MUCH faster than DP but *may*
+have configuration issues depending on your cluster.
+
+For a deeper understanding of what lightning is doing, feel free to read [this guide](https://medium.com/@_willfalcon/9-tips-for-training-lightning-fast-neural-networks-in-pytorch-8e63a502f565).   
+
 ---
 #### 16-bit mixed precision
 16 bit precision can cut your memory footprint by half. If using volta architecture GPUs it can give a dramatic training speed-up as well.    
@@ -37,16 +57,69 @@ Make sure you're on a GPU machine. You can set as many GPUs as you want.
 In this setting, the model will run on all 8 GPUs at once using DataParallel under the hood.
 ```python
 # set these flags
-os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
-os.environ["CUDA_VISIBLE_DEVICES"] = "0,1,2,3,4,5,6,7"
+# lightning sets these flags for you automatically
+# no need to set yourself
+# os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
+# os.environ["CUDA_VISIBLE_DEVICES"] = "0,1,2,3,4,5,6,7"
 
-# DEFAULT
-trainer = Trainer(gpus=[0,1,2,3,4,5,6,7])
+
+# to use DataParallel (default)
+trainer = Trainer(gpus=[0,1,2,3,4,5,6,7], distributed_backend='dp')
+
+# RECOMMENDED use DistributedDataParallel
+trainer = Trainer(gpus=[0,1,2,3,4,5,6,7], distributed_backend='ddp')
 ```
 
 ---
 #### Multi-node
-COMING SOON.
+Multi-node training is easily done by specifying these flags.
+```python
+# train on 12*8 GPUs
+trainer = Trainer(gpus=[0,1,2,3,4,5,6,7], nb_gpu_nodes=12)
+```
+
+In addition, make sure to set up your SLURM job correctly via the [SlurmClusterObject](https://williamfalcon.github.io/test-tube/hpc/SlurmCluster/). In particular, specify the number of tasks per node correctly.
+
+```python
+cluster = SlurmCluster(
+    hyperparam_optimizer=test_tube.HyperOptArgumentParser(),
+    log_path='/some/path/to/save',
+)
+
+# OPTIONAL FLAGS WHICH MAY BE CLUSTER DEPENDENT
+# which interface your nodes use for communication
+cluster.add_command('export NCCL_SOCKET_IFNAME=^docker0,lo')
+
+# see output of the NCCL connection process
+# NCCL is how the nodes talk to each other
+cluster.add_command('export NCCL_DEBUG=INFO')
+
+# setting a master port here is a good idea.
+cluster.add_command(f'export MASTER_PORT={PORT}')
+
+# good to load the latest NCCL version
+cluster.load_modules(['NCCL/2.4.7-1-cuda.10.0'])
+
+# configure cluster
+cluster.per_experiment_nb_nodes = 12 
+cluster.per_experiment_nb_gpus = 8
+
+cluster.add_slurm_cmd(cmd='ntasks-per-node', value=8, comment='1 task per gpu')
+```
+
+Finally, make sure to add a distributed sampler to your dataset. The distributed sampler copies a 
+portion of your dataset onto each GPU. (World_size = gpus_per_node * nb_nodes).   
+
+```python
+# ie: this:
+dataset = myDataset()
+dataloader = Dataloader(dataset)
+
+# becomes:
+dataset = myDataset()
+dist_sampler = torch.utils.data.distributed.DistributedSampler(dataset)
+dataloader = Dataloader(dataset, sampler=dist_sampler)
+```
 
 ---
 #### Self-balancing architecture

docs/Trainer/SLURM Managed Cluster.md

@@ -24,7 +24,7 @@ hparams = parser.parse_args()
 ```    
     
      
-(2). Define the cluster options (over 5 nodes and 8 gpus)    
+(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

docs/Trainer/Training Loop.md

@@ -19,7 +19,7 @@ Cut the learning rate by 10 at every epoch listed in this list.
 trainer = Trainer(lr_scheduler_milestones=None)
 
 # cut LR by 10 at 100, 200, and 300 epochs 
-trainer = Trainer(lr_scheduler_milestones=[100, 200, 300])
+trainer = Trainer(lr_scheduler_milestones='100, 200, 300')
 ```
 
 ---
@@ -38,6 +38,16 @@ Use this to turn off early stopping and run training to the [max_epoch](#force-t
 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.

docs/Trainer/index.md

@@ -24,10 +24,8 @@ But of course the fun is in all the advanced things it can do:
 
 **Computing cluster (SLURM)**    
 
-- Automatic checkpointing   
-- Automatic saving, loading  
-- Running grid search on a cluster
-- Walltime auto-resubmit   
+- [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**  
 

docs/doc_requirements.txt

@@ -0,0 +1 @@
+mkdocs-material==4.4.0

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

@@ -1,35 +1,33 @@
-# PYTORCH-LIGHTNING DOCUMENTATION
+###### New project Quick Start    
+To start a new project define these two files.   
 
-###### Main Docs
-- [LightningModule](Pytorch-Lightning/LightningModule)  
-- [Trainer](Trainer/)  
-
-###### New project Quick Start
-1. [Define a LightningModule](https://github.com/williamFalcon/pytorch-lightning/blob/master/examples/new_project_templates/lightning_module_template.py)  
+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   
-- Single GPU example   
-- Multi-gpu example 
-- SLURM cluster grid search example      
+- [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
-- Model loading 
+- [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)    
+###### Computing cluster (SLURM)    
 
-- Automatic checkpointing   
-- Automatic saving, loading  
-- Running grid search on a cluster 
-- Walltime auto-resubmit   
+- [Running grid search on a cluster](https://williamfalcon.github.io/pytorch-lightning/Trainer/SLURM%20Managed%20Cluster#running-grid-search-on-a-cluster) 
+- [Walltime auto-resubmit](https://williamfalcon.github.io/pytorch-lightning/Trainer/SLURM%20Managed%20Cluster#walltime-auto-resubmit)   
 
-######Debugging  
+###### Debugging  
 
 - [Fast dev run](https://williamfalcon.github.io/pytorch-lightning/Trainer/debugging/#fast-dev-run)
 - [Inspect gradient norms](https://williamfalcon.github.io/pytorch-lightning/Trainer/debugging/#inspect-gradient-norms)
@@ -39,7 +37,7 @@
 - [Pring which gradients are nan](https://williamfalcon.github.io/pytorch-lightning/Trainer/debugging/#print-which-gradients-are-nan)
 
 
-######Distributed training    
+###### Distributed training    
 
 - [16-bit mixed precision](https://williamfalcon.github.io/pytorch-lightning/Trainer/Distributed%20training/#16-bit-mixed-precision)
 - [Multi-GPU](https://williamfalcon.github.io/pytorch-lightning/Trainer/Distributed%20training/#Multi-GPU)
@@ -48,22 +46,23 @@
 - [Self-balancing architecture](https://williamfalcon.github.io/pytorch-lightning/Trainer/Distributed%20training/#self-balancing-architecture)
 
 
-######Experiment Logging   
+###### Experiment Logging   
 
 - [Display metrics in progress bar](https://williamfalcon.github.io/pytorch-lightning/Trainer/Logging/#display-metrics-in-progress-bar)
 - Log arbitrary metrics
 - [Log metric row every k batches](https://williamfalcon.github.io/pytorch-lightning/Trainer/Logging/#log-metric-row-every-k-batches)
-- [Process position](Logging/#process-position)
+- [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    
+###### 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)
 

docs/source/examples/__init__.py

@@ -1 +0,0 @@
-from .example_model import ExampleModel

docs/source/examples/basic_trainer.py

@@ -1,74 +0,0 @@
-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_function=None,
-        save_best_only=True,
-        verbose=True,
-        monitor='val_acc',
-        mode='min'
-    )
-
-    # configure trainer
-    trainer = Trainer(
-        experiment=exp,
-        checkpoint_callback=checkpoint,
-        early_stop_callback=early_stop,
-    )
-
-    # train model
-    trainer.fit(model)
-
-
-if __name__ == '__main__':
-
-    # use default args given by lightning
-    root_dir = os.path.split(os.path.dirname(sys.modules['__main__'].__file__))[0]
-    parent_parser = HyperOptArgumentParser(strategy='random_search', add_help=False)
-    add_default_args(parent_parser, root_dir)
-
-    # allow model to overwrite or extend args
-    parser = ExampleModel.add_model_specific_args(parent_parser)
-    hyperparams = parser.parse_args()
-
-    # train model
-    main(hyperparams)

docs/source/examples/fully_featured_trainer.py

@@ -1,210 +0,0 @@
-import os
-import sys
-import numpy as np
-from time import sleep
-import torch
-
-from test_tube import HyperOptArgumentParser, Experiment, SlurmCluster
-from pytorch_lightning.models.trainer import Trainer
-from pytorch_lightning.utils.arg_parse import add_default_args
-
-from pytorch_lightning.callbacks import EarlyStopping, ModelCheckpoint
-
-SEED = 2334
-torch.manual_seed(SEED)
-np.random.seed(SEED)
-
-# ---------------------
-# DEFINE MODEL HERE
-# ---------------------
-from docs.source.examples.example_model import ExampleModel
-# ---------------------
-
-AVAILABLE_MODELS = {
-    'model_template': ExampleModel
-}
-
-
-"""
-Allows training by using command line arguments
-Run by:
-# TYPE YOUR RUN COMMAND HERE
-"""
-
-
-def main_local(hparams):
-    main(hparams, None, None)
-
-
-def main(hparams, cluster, results_dict):
-    """
-    Main training routine specific for this project
-    :param hparams:
-    :return:
-    """
-    on_gpu = hparams.gpus is not None and torch.cuda.is_available()
-
-    device = 'cuda' if on_gpu else 'cpu'
-    hparams.__setattr__('device', device)
-    hparams.__setattr__('on_gpu', on_gpu)
-    hparams.__setattr__('nb_gpus', torch.cuda.device_count())
-    hparams.__setattr__('inference_mode', hparams.model_load_weights_path is not None)
-
-    # delay each training start to not overwrite logs
-    process_position, current_gpu = TRAINING_MODEL.get_process_position(hparams.gpus)
-    sleep(process_position + 1)
-
-    # init experiment
-    log_dir = os.path.dirname(os.path.realpath(__file__))
-    exp = Experiment(
-        name='test_tube_exp',
-        debug=True,
-        save_dir=log_dir,
-        version=0,
-        autosave=False,
-        description='test demo'
-    )
-
-    exp.argparse(hparams)
-    exp.save()
-
-    # build model
-    print('loading model...')
-    model = TRAINING_MODEL(hparams)
-    print('model built')
-
-    # callbacks
-    early_stop = EarlyStopping(
-        monitor=hparams.early_stop_metric,
-        patience=hparams.early_stop_patience,
-        verbose=True,
-        mode=hparams.early_stop_mode
-    )
-
-    model_save_path = '{}/{}/{}'.format(hparams.model_save_path, exp.name, exp.version)
-    checkpoint = ModelCheckpoint(
-        filepath=model_save_path,
-        save_function=None,
-        save_best_only=True,
-        verbose=True,
-        monitor=hparams.model_save_monitor_value,
-        mode=hparams.model_save_monitor_mode
-    )
-
-    # gpus are ; separated for inside a node and , within nodes
-    gpu_list = None
-    if hparams.gpus is not None:
-        gpu_list = [int(x) for x in hparams.gpus.split(';')]
-
-    # configure trainer
-    trainer = Trainer(
-        experiment=exp,
-        cluster=cluster,
-        checkpoint_callback=checkpoint,
-        early_stop_callback=early_stop,
-        gpus=gpu_list
-    )
-
-    # train model
-    trainer.fit(model)
-
-
-def get_default_parser(strategy, root_dir):
-
-    possible_model_names = list(AVAILABLE_MODELS.keys())
-    parser = HyperOptArgumentParser(strategy=strategy, add_help=False)
-    add_default_args(parser, root_dir, possible_model_names=possible_model_names, rand_seed=SEED)
-    return parser
-
-
-def get_model_name(args):
-    for i, arg in enumerate(args):
-        if 'model_name' in arg:
-            return args[i+1]
-
-
-def optimize_on_cluster(hyperparams):
-    # enable cluster training
-    cluster = SlurmCluster(
-        hyperparam_optimizer=hyperparams,
-        log_path=hyperparams.tt_save_path,
-        test_tube_exp_name=hyperparams.tt_name
-    )
-
-    # email for cluster coms
-    cluster.notify_job_status(email='add_email_here', on_done=True, on_fail=True)
-
-    # configure cluster
-    cluster.per_experiment_nb_gpus = hyperparams.per_experiment_nb_gpus
-    cluster.job_time = '48:00:00'
-    cluster.gpu_type = '1080ti'
-    cluster.memory_mb_per_node = 48000
-
-    # any modules for code to run in env
-    cluster.add_command('source activate pytorch_lightning')
-
-    # name of exp
-    job_display_name = hyperparams.tt_name.split('_')[0]
-    job_display_name = job_display_name[0:3]
-
-    # run hopt
-    print('submitting jobs...')
-    cluster.optimize_parallel_cluster_gpu(
-        main,
-        nb_trials=hyperparams.nb_hopt_trials,
-        job_name=job_display_name
-    )
-
-
-if __name__ == '__main__':
-
-    model_name = get_model_name(sys.argv)
-    if model_name is None:
-        model_name = 'model_template'
-
-    # use default args
-    root_dir = os.path.dirname(os.path.realpath(__file__))
-    parent_parser = get_default_parser(strategy='random_search', root_dir=root_dir)
-
-    # allow model to overwrite or extend args
-    TRAINING_MODEL = AVAILABLE_MODELS[model_name]
-    parser = TRAINING_MODEL.add_model_specific_args(parent_parser, root_dir)
-    hyperparams = parser.parse_args()
-
-    # format GPU layout
-    os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
-
-    # ---------------------
-    # RUN TRAINING
-    # ---------------------
-
-    # cluster and CPU
-    if hyperparams.on_cluster:
-        # run on HPC cluster
-        print('RUNNING ON SLURM CLUSTER')
-        gpu_ids = hyperparams.gpus.split(';')
-        os.environ["CUDA_VISIBLE_DEVICES"] = ','.join(gpu_ids)
-        optimize_on_cluster(hyperparams)
-
-    elif hyperparams.gpus is None:
-        # run on cpu
-        print('RUNNING ON CPU')
-        main(hyperparams, None, None)
-
-    # single or multiple GPUs on same machine
-    gpu_ids = hyperparams.gpus.split(';')
-    if hyperparams.interactive:
-        # run on 1 gpu
-        print(f'RUNNING INTERACTIVE MODE ON GPUS. gpu ids: {gpu_ids}')
-        os.environ["CUDA_VISIBLE_DEVICES"] = ','.join(gpu_ids)
-        main(hyperparams, None, None)
-
-    else:
-        # multiple GPUs on same machine
-        print(f'RUNNING MULTI GPU. GPU ids: {gpu_ids}')
-        hyperparams.optimize_parallel_gpu(
-            main_local,
-            gpu_ids=gpu_ids,
-            nb_trials=hyperparams.nb_hopt_trials,
-            nb_workers=len(gpu_ids)
-        )

examples/new_project_templates/trainer_gpu_cluster_template.py

@@ -1,210 +0,0 @@
-import os
-import sys
-import numpy as np
-from time import sleep
-import torch
-
-from test_tube import HyperOptArgumentParser, Experiment, SlurmCluster
-from pytorch_lightning.models.trainer import Trainer
-from pytorch_lightning.utils.arg_parse import add_default_args
-
-from pytorch_lightning.callbacks import EarlyStopping, ModelCheckpoint
-
-SEED = 2334
-torch.manual_seed(SEED)
-np.random.seed(SEED)
-
-# ---------------------
-# DEFINE MODEL HERE
-# ---------------------
-from examples.new_project_templates.lightning_module_template import LightningTemplateModel
-# ---------------------
-
-AVAILABLE_MODELS = {
-    'model_template': 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:
-    """
-    on_gpu = hparams.gpus is not None and torch.cuda.is_available()
-
-    device = 'cuda' if on_gpu else 'cpu'
-    hparams.__setattr__('device', device)
-    hparams.__setattr__('on_gpu', on_gpu)
-    hparams.__setattr__('nb_gpus', torch.cuda.device_count())
-    hparams.__setattr__('inference_mode', hparams.model_load_weights_path is not None)
-
-    # delay each training start to not overwrite logs
-    process_position, current_gpu = TRAINING_MODEL.get_process_position(hparams.gpus)
-    sleep(process_position + 1)
-
-    # init experiment
-    log_dir = os.path.dirname(os.path.realpath(__file__))
-    exp = Experiment(
-        name='test_tube_exp',
-        debug=True,
-        save_dir=log_dir,
-        version=0,
-        autosave=False,
-        description='test demo'
-    )
-
-    exp.argparse(hparams)
-    exp.save()
-
-    # build model
-    print('loading model...')
-    model = TRAINING_MODEL(hparams)
-    print('model built')
-
-    # callbacks
-    early_stop = EarlyStopping(
-        monitor=hparams.early_stop_metric,
-        patience=hparams.early_stop_patience,
-        verbose=True,
-        mode=hparams.early_stop_mode
-    )
-
-    model_save_path = '{}/{}/{}'.format(hparams.model_save_path, exp.name, exp.version)
-    checkpoint = ModelCheckpoint(
-        filepath=model_save_path,
-        save_function=None,
-        save_best_only=True,
-        verbose=True,
-        monitor=hparams.model_save_monitor_value,
-        mode=hparams.model_save_monitor_mode
-    )
-
-    # gpus are ; separated for inside a node and , within nodes
-    gpu_list = None
-    if hparams.gpus is not None:
-        gpu_list = [int(x) for x in hparams.gpus.split(';')]
-
-    # configure trainer
-    trainer = Trainer(
-        experiment=exp,
-        cluster=cluster,
-        checkpoint_callback=checkpoint,
-        early_stop_callback=early_stop,
-        gpus=gpu_list
-    )
-
-    # train model
-    trainer.fit(model)
-
-
-def get_default_parser(strategy, root_dir):
-
-    possible_model_names = list(AVAILABLE_MODELS.keys())
-    parser = HyperOptArgumentParser(strategy=strategy, add_help=False)
-    add_default_args(parser, root_dir, possible_model_names=possible_model_names, rand_seed=SEED)
-    return parser
-
-
-def get_model_name(args):
-    for i, arg in enumerate(args):
-        if 'model_name' in arg:
-            return args[i+1]
-
-
-def optimize_on_cluster(hyperparams):
-    # enable cluster training
-    cluster = SlurmCluster(
-        hyperparam_optimizer=hyperparams,
-        log_path=hyperparams.tt_save_path,
-        test_tube_exp_name=hyperparams.tt_name
-    )
-
-    # email for cluster coms
-    cluster.notify_job_status(email='add_email_here', on_done=True, on_fail=True)
-
-    # configure cluster
-    cluster.per_experiment_nb_gpus = hyperparams.per_experiment_nb_gpus
-    cluster.job_time = '48:00:00'
-    cluster.gpu_type = '1080ti'
-    cluster.memory_mb_per_node = 48000
-
-    # any modules for code to run in env
-    cluster.add_command('source activate pytorch_lightning')
-
-    # name of exp
-    job_display_name = hyperparams.tt_name.split('_')[0]
-    job_display_name = job_display_name[0:3]
-
-    # run hopt
-    print('submitting jobs...')
-    cluster.optimize_parallel_cluster_gpu(
-        main,
-        nb_trials=hyperparams.nb_hopt_trials,
-        job_name=job_display_name
-    )
-
-
-if __name__ == '__main__':
-
-    model_name = get_model_name(sys.argv)
-    if model_name is None:
-        model_name = 'model_template'
-
-    # use default args
-    root_dir = os.path.dirname(os.path.realpath(__file__))
-    parent_parser = get_default_parser(strategy='random_search', root_dir=root_dir)
-
-    # allow model to overwrite or extend args
-    TRAINING_MODEL = AVAILABLE_MODELS[model_name]
-    parser = TRAINING_MODEL.add_model_specific_args(parent_parser, root_dir)
-    hyperparams = parser.parse_args()
-
-    # format GPU layout
-    os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
-
-    # ---------------------
-    # RUN TRAINING
-    # ---------------------
-
-    # cluster and CPU
-    if hyperparams.on_cluster:
-        # run on HPC cluster
-        print('RUNNING ON SLURM CLUSTER')
-        gpu_ids = hyperparams.gpus.split(';')
-        os.environ["CUDA_VISIBLE_DEVICES"] = ','.join(gpu_ids)
-        optimize_on_cluster(hyperparams)
-
-    elif hyperparams.gpus is None:
-        # run on cpu
-        print('RUNNING ON CPU')
-        main(hyperparams, None, None)
-
-    # single or multiple GPUs on same machine
-    gpu_ids = hyperparams.gpus.split(';')
-    if hyperparams.interactive:
-        # run on 1 gpu
-        print(f'RUNNING INTERACTIVE MODE ON GPUS. gpu ids: {gpu_ids}')
-        os.environ["CUDA_VISIBLE_DEVICES"] = ','.join(gpu_ids)
-        main(hyperparams, None, None)
-
-    else:
-        # multiple GPUs on same machine
-        print(f'RUNNING MULTI GPU. GPU ids: {gpu_ids}')
-        hyperparams.optimize_parallel_gpu(
-            main_local,
-            gpu_ids=gpu_ids,
-            nb_trials=hyperparams.nb_hopt_trials,
-            nb_workers=len(gpu_ids)
-        )

mkdocs.yml

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

pytorch_lightning/__init__.py

@@ -1 +1,2 @@
-from .models import Trainer
+from .models import Trainer
+from .root_module.root_module import LightningModule

pytorch_lightning/callbacks/pt_callbacks.py

@@ -1,6 +1,6 @@
 import numpy as np
 import os, shutil
-from pytorch_lightning.pt_overrides.override_data_parallel import LightningDataParallel
+from pytorch_lightning.pt_overrides.override_data_parallel import LightningDistributedDataParallel
 
 
 class Callback(object):
@@ -34,7 +34,7 @@ class Callback(object):
         self.params = params
 
     def set_model(self, model):
-        if type(model) is LightningDataParallel:
+        if type(model) is LightningDistributedDataParallel:
             model = model.module
         self.model = model
 
@@ -170,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/examples/__init__.py

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

pytorch_lightning/examples/new_project_templates/lightning_module_template.py

@@ -7,6 +7,8 @@ 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
 
@@ -26,6 +28,9 @@ class LightningTemplateModel(LightningModule):
 
         self.batch_size = hparams.batch_size
 
+        # if you specify an example input, the summary will show input/output for each layer
+        self.example_input_array = torch.rand(5, 28 * 28)
+
         # build model
         self.__build_model()
 
@@ -76,15 +81,21 @@ class LightningTemplateModel(LightningModule):
         # forward pass
         x, y = data_batch
         x = x.view(x.size(0), -1)
+
         y_hat = self.forward(x)
 
         # calculate loss
         loss_val = self.loss(y, y_hat)
 
+        # in DP mode (default) make sure if result is scalar, there's another dim in the beginning
+        if self.trainer.use_dp:
+            loss_val = loss_val.unsqueeze(0)
+
         output = OrderedDict({
-            'loss': loss_val,
-            'tqdm_metrics': {}
+            'loss': loss_val
         })
+
+        # can also return just a scalar instead of a dict (return loss_val)
         return output
 
     def validation_step(self, data_batch, batch_i):
@@ -102,11 +113,22 @@ class LightningTemplateModel(LightningModule):
         # acc
         labels_hat = torch.argmax(y_hat, dim=1)
         val_acc = torch.sum(y == labels_hat).item() / (len(y) * 1.0)
+        val_acc = torch.tensor(val_acc)
+
+        if self.on_gpu:
+            val_acc = val_acc.cuda(loss_val.device.index)
+
+        # in DP mode (default) make sure if result is scalar, there's another dim in the beginning
+        if self.trainer.use_dp:
+            loss_val = loss_val.unsqueeze(0)
+            val_acc = val_acc.unsqueeze(0)
 
         output = OrderedDict({
             'val_loss': loss_val,
-            'val_acc': torch.tensor(val_acc),
+            'val_acc': val_acc,
         })
+
+        # can also return just a scalar instead of a dict (return loss_val)
         return output
 
     def validation_end(self, outputs):
@@ -115,6 +137,10 @@ class LightningTemplateModel(LightningModule):
         :param outputs: list of individual outputs of each validation step
         :return:
         """
+        # if returned a scalar from validation_step, outputs is a list of tensor scalars
+        # we return just the average in this case (if we want)
+        # return torch.stack(outputs).mean()
+
         val_loss_mean = 0
         val_acc_mean = 0
         for output in outputs:
@@ -126,9 +152,6 @@ class LightningTemplateModel(LightningModule):
         tqdm_dic = {'val_loss': val_loss_mean.item(), 'val_acc': val_acc_mean.item()}
         return tqdm_dic
 
-    def update_tng_log_metrics(self, logs):
-        return logs
-
     # ---------------------
     # MODEL SAVING
     # ---------------------
@@ -154,13 +177,25 @@ class LightningTemplateModel(LightningModule):
     def __dataloader(self, train):
         # init data generators
         transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,), (1.0,))])
-
         dataset = MNIST(root=self.hparams.data_root, train=train, transform=transform, download=True)
 
-        loader = torch.utils.data.DataLoader(
+        # when using multi-node we need to add the datasampler
+        train_sampler = None
+        batch_size = self.hparams.batch_size
+
+        try:
+            if self.on_gpu:
+                train_sampler = DistributedSampler(dataset, rank=self.trainer.proc_rank)
+                batch_size = batch_size // self.trainer.world_size  # scale batch size
+        except Exception as e:
+            pass
+
+        should_shuffle = train_sampler is None
+        loader = DataLoader(
             dataset=dataset,
-            batch_size=self.hparams.batch_size,
-            shuffle=True
+            batch_size=batch_size,
+            shuffle=should_shuffle,
+            sampler=train_sampler
         )
 
         return loader
@@ -196,7 +231,7 @@ class LightningTemplateModel(LightningModule):
         return self._test_dataloader
 
     @staticmethod
-    def add_model_specific_args(parent_parser, root_dir):
+    def add_model_specific_args(parent_parser, root_dir): # pragma: no cover
         """
         Parameters you define here will be available to your model through self.hparams
         :param parent_parser:
@@ -210,16 +245,19 @@ class LightningTemplateModel(LightningModule):
 
         # 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
+        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, type=float, options=[0.0001, 0.0005, 0.001, 0.005],
+        parser.opt_list('--learning_rate', default=0.001*8, type=float, options=[0.0001, 0.0005, 0.001, 0.005],
                         tunable=False)
-        parser.opt_list('--batch_size', default=256, type=int, options=[32, 64, 128, 256], tunable=False)
         parser.opt_list('--optimizer_name', default='adam', type=str, options=['adam'], tunable=False)
+
+        # if using 2 nodes with 4 gpus each the batch size here (256) will be 256 / (2*8) = 16 per gpu
+        parser.opt_list('--batch_size', default=256*8, type=int, options=[32, 64, 128, 256], tunable=False,
+                        help='batch size will be divided over all the gpus being used across all nodes')
         return parser

pytorch_lightning/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)

pytorch_lightning/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)

pytorch_lightning/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)

pytorch_lightning/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)

pytorch_lightning/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)

pytorch_lightning/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)

pytorch_lightning/examples/new_project_templates/trainer_cpu_template.py

@@ -41,7 +41,6 @@ def main(hparams):
     model_save_path = '{}/{}/{}'.format(hparams.model_save_path, exp.name, exp.version)
     checkpoint = ModelCheckpoint(
         filepath=model_save_path,
-        save_function=None,
         save_best_only=True,
         verbose=True,
         monitor='val_acc',

pytorch_lightning/models/trainer.py

@@ -1,21 +1,42 @@
-import torch
-import tqdm
-import numpy as np
-from pytorch_lightning.root_module.memory import get_gpu_memory_map
+"""
+The trainer handles all the logic for running a val loop, training loop, distributing, etc...
+"""
+import subprocess
 import traceback
-from pytorch_lightning.root_module.model_saving import TrainerIO
-from torch.optim.lr_scheduler import MultiStepLR
-from pytorch_lightning.pt_overrides.override_data_parallel import LightningDataParallel
+import warnings
+import os
 import pdb
+import re
+
+import torch
+from torch.utils.data.distributed import DistributedSampler
+from torch.optim.lr_scheduler import MultiStepLR
+import torch.multiprocessing as mp
+import torch.distributed as dist
+import numpy as np
+import tqdm
+
+from pytorch_lightning.root_module.memory import get_gpu_memory_map
+from pytorch_lightning.root_module.model_saving import TrainerIO
+from pytorch_lightning.pt_overrides.override_data_parallel import LightningDistributedDataParallel, LightningDataParallel
+from pytorch_lightning.utils.debugging import MisconfigurationException
 
 try:
     from apex import amp
     APEX_AVAILABLE = True
-except ModuleNotFoundError:
+except ModuleNotFoundError: # pragma: no cover
     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):
@@ -32,10 +53,13 @@ class Trainer(TrainerIO):
 
     def __init__(self,
                  experiment,
-                 checkpoint_callback, early_stop_callback,
+                 early_stop_callback=None,
+                 checkpoint_callback=None,
+                 gradient_clip=0,
                  cluster=None,
                  process_position=0,
                  current_gpu_name=0,
+                 nb_gpu_nodes=1,
                  gpus=None,
                  progress_bar=True,
                  overfit_pct=0.0,
@@ -43,18 +67,58 @@ class Trainer(TrainerIO):
                  check_val_every_n_epoch=1,
                  fast_dev_run=False,
                  accumulate_grad_batches=1,
-                 enable_early_stop=True, max_nb_epochs=1000, min_nb_epochs=1,
-                 train_percent_check=1.0, val_percent_check=1.0, test_percent_check=1.0, val_check_interval=0.95,
+                 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 = gpus is not None and torch.cuda.is_available()
@@ -64,8 +128,12 @@ class Trainer(TrainerIO):
         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
@@ -77,8 +145,73 @@ class Trainer(TrainerIO):
         self.lr_schedulers = []
         self.amp_level = amp_level
         self.print_nan_grads = print_nan_grads
-        self.data_parallel_device_ids = gpus
-        self.data_parallel = gpus is not None and len(gpus) > 0
+        self.data_parallel_device_ids = None
+        self.world_size = 1
+        self.node_rank = 0
+        self.use_ddp = False
+        self.use_dp = False
+
+        # training bookeeping
+        self.total_batch_nb = 0
+        self.running_loss = []
+        self.avg_loss = 0
+        self.batch_nb = 0
+        self.tqdm_metrics = {}
+        self.nb_val_batches = None
+        self.nb_tng_batches = None
+        self.nb_test_batches = None
+
+
+        # gpus come in as a string.
+        # if gpus = -1 then use all available devices
+        # otherwise, split the string using commas
+        if gpus is not None:
+            if type(gpus) is list:
+                self.data_parallel_device_ids = gpus
+            elif type(gpus) is str:
+                if gpus == '-1':
+                    self.data_parallel_device_ids = list(range(0, torch.cuda.device_count()))
+                else:
+                    self.data_parallel_device_ids = [int(x.strip()) for x in gpus.split(',')]
+            else:
+                raise Exception('gpus has to be a string or list of ids')
+
+            # set the correct cuda visible devices (using pci order)
+            os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
+            os.environ["CUDA_VISIBLE_DEVICES"] = ','.join([str(x) for x in self.data_parallel_device_ids])
+            print(f'VISIBLE GPUS: {os.environ["CUDA_VISIBLE_DEVICES"]}')
+
+        # make DP and DDP mutually exclusive
+        # single GPU will also use DP with devices=[0]
+        have_gpus = self.data_parallel_device_ids is not None and len(self.data_parallel_device_ids) > 0
+        if have_gpus:
+            self.use_dp = distributed_backend == 'dp'
+            self.use_ddp = distributed_backend == 'ddp'
+
+            # use ddp automatically if nb_gpu_nodes > 1
+            if nb_gpu_nodes > 1 and self.use_dp:  # pragma: no cover
+                self.use_ddp = True
+                self.use_dp = False
+                w = 'DataParallel does not support nb_gpu_nodes > 1. ' \
+                    'Switching to DistributedDataParallel for you. ' \
+                    'To silence this warning set distributed_backend=ddp'
+                warnings.warn(w)
+
+        # extract SLURM flag vars
+        # whenever we have the correct number of tasks, we let slurm manage processes
+        # otherwise we launch the required number of processes
+        if self.use_ddp:
+            self.nb_requested_gpus = len(self.data_parallel_device_ids) * self.nb_gpu_nodes
+            self.nb_slurm_tasks = 0
+            try:
+                self.nb_slurm_tasks = int(os.environ['SLURM_NTASKS'])
+                self.is_slurm_managing_tasks = self.nb_slurm_tasks == self.nb_requested_gpus
+            except Exception as e:
+                # likely not on slurm, so set the slurm managed flag to false
+                self.is_slurm_managing_tasks = False
+
+        # process info
+        self.proc_rank = 0
 
         # training state
         self.optimizers = None
@@ -101,11 +234,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))
 
-        # apex test
+        # 16 bit mixed precision training using apex
         self.use_amp = use_amp and APEX_AVAILABLE
         if self.use_amp:
             print('using 16bit precision')
 
+        if use_amp and not APEX_AVAILABLE:  # pragma: no cover
+            msg = '''
+            You set use_amp=True but do not have apex installed.
+            Install apex first using this guide and rerun with use_amp=True: 
+            https://github.com/NVIDIA/apex#linux
+            
+            this run will NOT use 16 bit precision
+            '''
+            raise ModuleNotFoundError(msg)
+
+    @property
+    def data_parallel(self):
+        return self.use_dp or self.use_ddp
+
     def __determine_data_use_amount(self, train_percent_check, val_percent_check, test_percent_check, overfit_pct):
         """
         Use less data for debugging purposes
@@ -118,43 +265,52 @@ 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.model, f_name, None)
+        model = self.__get_model()
+        f_op = getattr(model, f_name, None)
         return callable(f_op)
 
     @property
     def __tng_tqdm_dic(self):
+        # ForkedPdb().set_trace()
         tqdm_dic = {
             'tng_loss': '{0:.3f}'.format(self.avg_loss),
-            '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, model):
-        # training bookeeping
-        self.total_batch_nb = 0
-        self.running_loss = []
-        self.avg_loss = 0
-        self.batch_nb = 0
-        self.tqdm_metrics = {}
+    @property
+    def tng_tqdm_dic(self):
+        """
+        Read-only for tqdm metrics
+        :return:
+        """
+        return self.__tng_tqdm_dic
+
+    def __layout_bookeeping(self):
 
         # determine number of training batches
-        self.nb_tng_batches = model.nb_batches(self.tng_dataloader)
+        self.nb_tng_batches = len(self.tng_dataloader)
         self.nb_tng_batches = int(self.nb_tng_batches * self.train_percent_check)
 
         # determine number of validation batches
-        self.nb_val_batches = model.nb_batches(self.val_dataloader)
+        self.nb_val_batches = len(self.val_dataloader)
         self.nb_val_batches = int(self.nb_val_batches * self.val_percent_check)
         self.nb_val_batches = max(1, self.nb_val_batches)
         self.nb_val_batches = self.nb_val_batches
 
         # determine number of test batches
-        self.nb_test_batches = model.nb_batches(self.test_dataloader)
+        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
@@ -175,12 +331,9 @@ class Trainer(TrainerIO):
         :param max_batches: Scalar
         :return:
         """
-        print('validating...')
-
         # enable eval mode
         model.zero_grad()
         model.eval()
-        model.from_lightning = True
 
         # disable gradients to save memory
         torch.set_grad_enabled(False)
@@ -191,7 +344,7 @@ class Trainer(TrainerIO):
         # run training
         for batch_i, data_batch in enumerate(dataloader):
 
-            if data_batch is None:
+            if data_batch is None:  # pragma: no cover
                 continue
 
             # stop short when on fast dev run
@@ -201,9 +354,12 @@ class Trainer(TrainerIO):
             # -----------------
             # RUN VALIDATION STEP
             # -----------------
-            if self.data_parallel:
+            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)
 
@@ -227,7 +383,7 @@ class Trainer(TrainerIO):
 
         return val_results
 
-    def __get_dataloaders(self, model):
+    def get_dataloaders(self, model):
         """
         Dataloaders are provided by the model
         :param model:
@@ -237,30 +393,199 @@ 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 MisconfigurationException(msg)
+
     # -----------------------------
     # MODEL TRAINING
     # -----------------------------
     def fit(self, model):
 
-        model.trainer = self
+        # 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()
 
-        # transfer data loaders from model
-        self.__get_dataloaders(model)
+            if self.is_slurm_managing_tasks:
+                task = int(os.environ['SLURM_LOCALID'])
+                self.ddp_train(task, model)
+            else:
+                msg = f"""
+                You requested {self.nb_requested_gpus} GPUs but launched {self.nb_slurm_tasks} slurm tasks. 
+                We will launch {self.nb_requested_gpus} processes for you. 
+                We recommend you let slurm manage the processes by setting: --ntasks-per-node={self.nb_requested_gpus}
+                If you're not using SLURM, ignore this message!
+                """
+                warnings.warn(msg)
+                mp.spawn(self.ddp_train, nprocs=len(self.data_parallel_device_ids), args=(model, ))
 
-        # init training constants
-        self.__layout_bookeeping(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:
+                raise MisconfigurationException('amp + cpu is not supported. Please use a GPU option')
+
+            self.__run_pretrain_routine(model)
+
+        # return 1 when finished
+        # used for testing or when we need to know that training succeeded
+        return 1
+
+    def __dp_train(self, model):
 
         # CHOOSE OPTIMIZER
         # filter out the weights that were done on gpu so we can load on good old cpus
         self.optimizers = model.configure_optimizers()
 
+        model.cuda(self.data_parallel_device_ids[0])
+
+        # check for this bug (amp + dp + !01 doesn't work)
+        # https://github.com/NVIDIA/apex/issues/227
+        if self.use_dp and self.use_amp:
+            m = f'amp level {self.amp_level} with DataParallel is not supported. ' \
+                f'See this note from NVIDIA for more info: https://github.com/NVIDIA/apex/issues/227. ' \
+                f'We recommend you switch to ddp if you want to use amp'
+            raise MisconfigurationException(m)
+
+        model = LightningDataParallel(model, device_ids=self.data_parallel_device_ids)
+
+        self.__run_pretrain_routine(model)
+
+    def ddp_train(self, gpu_nb, model):
+        """
+        Entry point into a DP thread
+        :param gpu_nb:
+        :param model:
+        :param cluster_obj:
+        :return:
+        """
+        # node rank using relative slurm id
+        # otherwise default to node rank 0
+        try:
+            node_id = os.environ['SLURM_NODEID']
+            self.node_rank = int(node_id)
+        except Exception as e:
+            self.node_rank = 0
+
+        # recover original exp before went into process
+        # init in write mode only on proc 0
+        self.experiment.debug = self.proc_rank > 0
+        self.experiment = self.experiment.get_non_ddp_exp()
+
+        # show progbar only on prog_rank 0
+        self.prog_bar = self.prog_bar and self.node_rank == 0 and gpu_nb == 0
+
+        # determine which process we are and world size
+        self.proc_rank = self.node_rank * len(self.data_parallel_device_ids) + gpu_nb
+        self.world_size = self.nb_gpu_nodes * len(self.data_parallel_device_ids)
+
+        # 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, optimizer = amp.initialize(
-                model, self.optimizers[0], opt_level=self.amp_level,
+            model, optimizers = amp.initialize(
+                model, self.optimizers, opt_level=self.amp_level,
             )
-            self.optimizers[0] = optimizer
-            model.trainer = self
+            self.optimizers = optimizers
+
+        model = LightningDistributedDataParallel(model, device_ids=[gpu_nb], find_unused_parameters=True)
+
+        # continue training routine
+        self.__run_pretrain_routine(model)
+
+    def __init_tcp_connection(self):
+        """
+        Connect all procs in the world using the env:// init
+        Use the first node as the root address
+        :param port:
+        :param tries:
+        :return:
+        """
+        # sets the appropriate port
+        try:
+            port = os.environ['MASTER_PORT']
+        except Exception as e:
+            port = 12910
+            os.environ['MASTER_PORT'] = f'{port}'
+
+        # figure out the root node addr
+        root_node = os.environ['SLURM_NODELIST'].split(' ')[0]
+        root_node = self.resolve_root_node_address(root_node)
+        os.environ['MASTER_ADDR'] = root_node
+
+        dist.init_process_group("nccl", rank=self.proc_rank, world_size=self.world_size)
+
+    def resolve_root_node_address(self, root_node):
+        try:
+            if '[' in root_node:
+                name = root_node.split('[')[0]
+                number = root_node.split(',')[0]
+                if '-' in number:
+                    number = number.split('-')[0]
+
+                number = re.sub('[^0-9]', '', number)
+                root_node = name + number
+
+        except Exception as e:
+            root_node = '127.0.0.2'
+
+        return root_node
+
+    def __run_pretrain_routine(self, model):
+        """
+        Sanity check a few things before starting actual training
+        :param model:
+        :return:
+        """
+        ref_model = model
+        if self.data_parallel:
+            ref_model = model.module
+
+        ref_model.trainer = self
+
+        # set local properties on the model
+        ref_model.on_gpu = self.on_gpu
+
+        # transfer data loaders from model
+        self.get_dataloaders(ref_model)
+
+        # init training constants
+        self.__layout_bookeeping()
 
         # add lr schedulers
         if self.lr_scheduler_milestones is not None:
@@ -269,20 +594,22 @@ 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 = LightningDataParallel(model, device_ids=self.data_parallel_device_ids)
+        # give model convenience properties
+        ref_model.trainer = self
+        ref_model.experiment = self.experiment
 
         # run tiny validation to make sure program won't crash during val
         _ = self.validate(model, self.val_dataloader, max_batches=self.nb_sanity_val_steps)
 
         # save exp to get started
-        self.experiment.save()
+        if self.proc_rank == 0:
+            self.experiment.save()
 
         # enable cluster checkpointing
-        if self.cluster is not None:
+        if self.cluster is not None:  # pragma: no cover
             self.enable_auto_hpc_walltime_manager()
 
         # ---------------------------
@@ -298,12 +625,12 @@ class Trainer(TrainerIO):
             for lr_scheduler in self.lr_schedulers:
                 lr_scheduler.step()
 
-            model = self.model.module if self.data_parallel else self.model
+            model = self.__get_model()
             model.current_epoch = epoch_nb
 
             # hook
             if self.__is_function_implemented('on_epoch_start'):
-                model = self.model.module if self.data_parallel else self.model
+                model = self.__get_model()
                 model.on_epoch_start()
 
             self.current_epoch = epoch_nb
@@ -318,7 +645,7 @@ class Trainer(TrainerIO):
                 self.batch_nb = batch_nb
                 self.global_step += 1
 
-                model = self.model.module if self.data_parallel else self.model
+                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
@@ -342,17 +669,16 @@ class Trainer(TrainerIO):
 
                 # when batch should be saved
                 if (batch_nb + 1) % self.log_save_interval == 0 or early_stop_epoch:
-                    self.experiment.save()
+                    if self.proc_rank == 0:
+                        self.experiment.save()
 
                 # when metrics should be logged
                 if batch_nb % self.add_log_row_interval == 0 or early_stop_epoch:
                     # count items in memory
                     # nb_params, nb_tensors = count_mem_items()
 
-                    if self.data_parallel:
-                        metrics = self.model.module.update_tng_log_metrics(self.__tng_tqdm_dic)
-                    else:
-                        metrics = self.model.update_tng_log_metrics(self.__tng_tqdm_dic)
+                    model = self.__get_model()
+                    metrics = self.__tng_tqdm_dic
 
                     # add gpu memory
                     if self.on_gpu:
@@ -361,18 +687,22 @@ class Trainer(TrainerIO):
 
                     # add norms
                     if self.track_grad_norm > 0:
-                        model = self.model.module if self.data_parallel else self.model
+                        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'):
-                    model = self.model.module if self.data_parallel else self.model
+                    model = self.__get_model()
                     model.on_batch_end()
 
                 # end epoch early
@@ -381,19 +711,37 @@ class Trainer(TrainerIO):
 
             # hook
             if self.__is_function_implemented('on_epoch_end'):
-                model = self.model.module if self.data_parallel else self.model
+                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 __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:
@@ -401,8 +749,8 @@ class Trainer(TrainerIO):
 
         # hook
         if self.__is_function_implemented('on_batch_start'):
-            model = self.model.module if self.data_parallel else self.model
-            response = model.on_batch_start(data_batch)
+            model_ref = self.__get_model()
+            response = model_ref.on_batch_start(data_batch)
 
             if response == -1:
                 return -1
@@ -412,39 +760,68 @@ class Trainer(TrainerIO):
 
         # forward pass
         # return a scalar value and a dic with tqdm metrics
-        if self.data_parallel:
+        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)
 
-        model_specific_tqdm_metrics_dic = output['tqdm_metrics']
-        loss = output['loss']
+        try:
+            model_specific_tqdm_metrics_dic = output['tqdm_metrics']
+        except Exception as e:
+            model_specific_tqdm_metrics_dic = {}
+
+        # if output dict doesn't have the keyword loss
+        # then assume the output=loss if scalar
+        try:
+            loss = output['loss']
+        except Exception as e:
+            if type(output) is torch.Tensor:
+                loss = output
 
         self.__add_tqdm_metrics(model_specific_tqdm_metrics_dic)
 
         # backward pass
         if self.use_amp:
+            # scale loss when using amp
             for optimizer in self.optimizers:
                 with amp.scale_loss(loss, optimizer) as scaled_loss:
                     scaled_loss.backward()
         else:
             loss.backward()
 
+        # insert after step hook
+        if self.__is_function_implemented('on_after_backward'):
+            model_ref = self.__get_model()
+            response = model_ref.on_after_backward()
+
         if self.print_nan_grads:
-            model = self.model.module if self.data_parallel else self.model
+            model = self.__get_model()
             for param in model.parameters():
                 print(param.grad.float().sum())
 
+        # avoid memory leaks
         self.batch_loss_value += loss.item()
 
         # gradient update with accumulated gradients
         if (self.batch_nb + 1) % self.accumulate_grad_batches == 0:
 
+            # clip gradients
+            if self.gradient_clip > 0:
+                model = self.__get_model()
+                torch.nn.utils.clip_grad_norm(model.parameters(), self.gradient_clip)
+
             # update gradients across all optimizers
             for optimizer in self.optimizers:
                 optimizer.step()
 
+                # insert after step hook
+                if self.__is_function_implemented('on_before_zero_grad'):
+                    model_ref = self.__get_model()
+                    response = model_ref.on_before_zero_grad(optimizer)
+
                 # clear gradients
                 optimizer.zero_grad()
 
@@ -464,7 +841,8 @@ class Trainer(TrainerIO):
 
         # 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
 
@@ -479,7 +857,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
@@ -493,7 +872,8 @@ 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)
@@ -505,5 +885,6 @@ class Trainer(TrainerIO):
             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

@@ -1,4 +1,7 @@
 from torch.nn import DataParallel
+from torch.nn.parallel import DistributedDataParallel
+import itertools
+from itertools import chain
 
 import threading
 import torch
@@ -6,7 +9,20 @@ from torch.cuda._utils import _get_device_index
 import pdb
 
 
-def get_a_var(obj):
+def _find_tensors(obj):  # pragma: no cover
+    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):  # pragma: no cover
     if isinstance(obj, torch.Tensor):
         return obj
 
@@ -26,11 +42,77 @@ 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)])
 
 
-def parallel_apply(modules, inputs, kwargs_tup=None, devices=None):
+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):  # pragma: no cover
+        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):  # pragma: no cover
     r"""Applies each `module` in :attr:`modules` in parallel on arguments
     contained in :attr:`inputs` (positional) and :attr:`kwargs_tup` (keyword)
     on each of :attr:`devices`.
@@ -102,4 +184,4 @@ def parallel_apply(modules, inputs, kwargs_tup=None, devices=None):
         if isinstance(output, Exception):
             raise output
         outputs.append(output)
-    return outputs
+    return outputs

pytorch_lightning/root_module/grads.py

@@ -27,14 +27,3 @@ class GradInformation(nn.Module):
         results['grad_{}_norm_total'.format(norm_type)] = round(total_norm.data.cpu().numpy().flatten()[0], 3)
         return results
 
-
-    def describe_grads(self):
-        for p in self.parameters():
-            g = p.grad.data.numpy().flatten()
-            print(np.max(g), np.min(g), np.mean(g))
-
-
-    def describe_params(self):
-        for p in self.parameters():
-            g = p.data.numpy().flatten()
-            print(np.max(g), np.min(g), np.mean(g))

pytorch_lightning/root_module/hooks.py

@@ -19,3 +19,27 @@ class ModelHooks(torch.nn.Module):
     def on_post_performance_check(self):
         pass
 
+    def on_tng_metrics(self, metrics):
+        pass
+
+    def on_before_zero_grad(self, optimizer):
+        """
+        Called after optimizer.step() and before optimizer.zero_grad()
+
+        for optimizer in optimizers:
+            optimizer.step()
+            model.on_before_zero_grad(optimizer) # < ---- called here
+            optimizer.zero_grad
+
+        :param optimizer:
+        :return:
+        """
+        pass
+
+    def on_after_backward(self):
+        """
+        Called after loss.backward() and before optimizers do anything
+        :return:
+        """
+        pass
+

pytorch_lightning/root_module/memory.py

@@ -33,33 +33,42 @@ class ModelSummary(object):
         mods = list(self.model.modules())
         in_sizes = []
         out_sizes = []
-        input_ = self.example_input_array
-        for i in range(1, len(mods)):
-            m = mods[i]
-            if type(input_) is list or type(input_) is tuple:
-                out = m(*input_)
-            else:
-                out = m(input_)
+        input_ = self.model.example_input_array
 
-            if type(input_) is tuple or type(input_) is list:
-                in_size = []
-                for x in input_:
-                    if type(x) is list:
-                        in_size.append(len(x))
-                    else:
-                        in_size.append(x.size())
-            else:
-                in_size = np.array(input_.size())
+        if self.model.on_gpu:
+            input_ = input_.cuda(0)
 
-            in_sizes.append(in_size)
+        if self.model.trainer.use_amp:
+            input_ = input_.half()
 
-            if type(out) is tuple or type(out) is list:
-                out_size = np.asarray([x.size() for x in out])
-            else:
-                out_size = np.array(out.size())
+        with torch.no_grad():
 
-            out_sizes.append(out_size)
-            input_ = out
+            for i in range(1, len(mods)):
+                m = mods[i]
+                if type(input_) is list or type(input_) is tuple:  # pragma: no cover
+                    out = m(*input_)
+                else:
+                    out = m(input_)
+
+                if type(input_) is tuple or type(input_) is list:  # pragma: no cover
+                    in_size = []
+                    for x in input_:
+                        if type(x) is list:
+                            in_size.append(len(x))
+                        else:
+                            in_size.append(x.size())
+                else:
+                    in_size = np.array(input_.size())
+
+                in_sizes.append(in_size)
+
+                if type(out) is tuple or type(out) is list:  # pragma: no cover
+                    out_size = np.asarray([x.size() for x in out])
+                else:
+                    out_size = np.array(out.size())
+
+                out_sizes.append(out_size)
+                input_ = out
 
         self.in_sizes = in_sizes
         self.out_sizes = out_sizes
@@ -114,13 +123,22 @@ class ModelSummary(object):
         Layer Name, Layer Type, Input Size, Output Size, Number of Parameters
         '''
 
-        df = pd.DataFrame( np.zeros( (len(self.layer_names), 3) ) )
-        df.columns = ['Name', 'Type', 'Params']
+        cols = ['Name', 'Type', 'Params']
+        if self.model.example_input_array is not None:
+            cols.extend(['In_sizes', 'Out_sizes'])
+
+        df = pd.DataFrame(np.zeros( (len(self.layer_names), len(cols))))
+        df.columns = cols
 
         df['Name'] = self.layer_names
         df['Type'] = self.layer_types
         df['Params'] = self.param_nums
 
+        if self.model.example_input_array is not None:
+
+            df['In_sizes'] = self.in_sizes
+            df['Out_sizes'] = self.out_sizes
+
         self.summary = df
         return
 
@@ -128,10 +146,13 @@ class ModelSummary(object):
         self.get_layer_names()
         self.get_parameter_sizes()
         self.get_parameter_nums()
+
+        if self.model.example_input_array is not None:
+            self.get_variable_sizes()
         self.make_summary()
 
 
-def print_mem_stack():
+def print_mem_stack(): # pragma: no cover
     for obj in gc.get_objects():
         try:
             if torch.is_tensor(obj) or (hasattr(obj, 'data') and torch.is_tensor(obj.data)):
@@ -140,7 +161,7 @@ def print_mem_stack():
             pass
 
 
-def count_mem_items():
+def count_mem_items(): # pragma: no cover
     nb_params = 0
     nb_tensors = 0
     for obj in gc.get_objects():

pytorch_lightning/root_module/model_saving.py

@@ -2,7 +2,7 @@ import torch
 import os
 import re
 import pdb
-from pytorch_lightning.pt_overrides.override_data_parallel import LightningDataParallel
+from pytorch_lightning.pt_overrides.override_data_parallel import LightningDistributedDataParallel, LightningDataParallel
 
 class ModelIO(object):
 
@@ -21,9 +21,32 @@ 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):
 
+    def __get_model(self):
+        print(type(self.model))
+        is_dp_module = type(self.model) is LightningDistributedDataParallel or type(self.model) is LightningDataParallel
+        model = self.model.module if is_dp_module else self.model
+        return model
+
     # --------------------
     # MODEL SAVE CHECKPOINT
     # --------------------
@@ -34,14 +57,19 @@ class TrainerIO(object):
         torch.save(checkpoint, filepath)
 
     def dump_checkpoint(self):
+
         checkpoint = {
             'epoch': self.current_epoch,
-            'checkpoint_callback_best': self.checkpoint_callback.best,
-            'early_stop_callback_wait': self.early_stop_callback.wait,
-            'early_stop_callback_patience': self.early_stop_callback.patience,
             'global_step': self.global_step
         }
 
+        if self.checkpoint_callback is not None:
+            checkpoint['checkpoint_callback_best'] = self.checkpoint_callback.best
+
+        if self.early_stop_callback is not None:
+            checkpoint['early_stop_callback_wait'] = self.early_stop_callback.wait
+            checkpoint['early_stop_callback_patience'] = self.early_stop_callback.patience
+
         optimizer_states = []
         for i, optimizer in enumerate(self.optimizers):
             optimizer_states.append(optimizer.state_dict())
@@ -49,7 +77,7 @@ class TrainerIO(object):
         checkpoint['optimizer_states'] = optimizer_states
 
         # request what to save from the model
-        model = self.model.module if type(self.model) is LightningDataParallel else self.model
+        model = self.__get_model()
         checkpoint_dict = model.get_save_dict()
 
         # merge trainer and model saving items
@@ -59,7 +87,7 @@ class TrainerIO(object):
     # --------------------
     # HPC IO
     # --------------------
-    def enable_auto_hpc_walltime_manager(self):
+    def enable_auto_hpc_walltime_manager(self):  # pragma: no cover
         if self.cluster is None:
             return
 
@@ -86,9 +114,13 @@ class TrainerIO(object):
         :param checkpoint:
         :return:
         """
-        self.checkpoint_callback.best = checkpoint['checkpoint_callback_best']
-        self.early_stop_callback.wait = checkpoint['early_stop_callback_wait']
-        self.early_stop_callback.patience = checkpoint['early_stop_callback_patience']
+        if self.checkpoint_callback is not None:
+            self.checkpoint_callback.best = checkpoint['checkpoint_callback_best']
+
+        if self.early_stop_callback is not None:
+            self.early_stop_callback.wait = checkpoint['early_stop_callback_wait']
+            self.early_stop_callback.patience = checkpoint['early_stop_callback_patience']
+
         self.global_step = checkpoint['global_step']
         self.current_epoch = checkpoint['epoch']
 
@@ -107,12 +139,19 @@ class TrainerIO(object):
         # save exp to make sure we get all the metrics
         experiment.save()
 
+        # close experiment to avoid issues
+        experiment.close()
+
         ckpt_number = self.max_ckpt_in_folder(folderpath) + 1
 
         if not os.path.exists(folderpath):
             os.makedirs(folderpath, exist_ok=True)
         filepath = '{}/hpc_ckpt_{}.ckpt'.format(folderpath, ckpt_number)
 
+        # give model a chance to do something on hpc_save
+        model = self.__get_model()
+        model.on_hpc_save()
+
         # request what to save from the model
         checkpoint_dict = self.dump_checkpoint()
 
@@ -131,9 +170,12 @@ class TrainerIO(object):
         self.restore_training_state(checkpoint)
 
         # load model state
-        model = self.model.module if type(self.model) is LightningDataParallel else self.model
+        model = self.__get_model()
         model.load_model_specific(checkpoint)
 
+        # call model hook
+        model.on_hpc_load()
+
     def max_ckpt_in_folder(self, path):
         files = os.listdir(path)
         files = [x for x in files if 'ckpt_' in x]

pytorch_lightning/root_module/optimization.py

@@ -1,22 +0,0 @@
-from torch import nn
-from torch import optim
-
-
-class OptimizerConfig(nn.Module):
-
-    def choose_optimizer(self, optimizer, params, optimizer_params, opt_name_key):
-        if optimizer == 'adam':
-            optimizer = optim.Adam(params, **optimizer_params)
-        if optimizer == 'sparse_adam':
-            optimizer = optim.SparseAdam(params, **optimizer_params)
-        if optimizer == 'sgd':
-            optimizer = optim.SGD(params, **optimizer_params)
-        if optimizer == 'adadelta':
-            optimizer = optim.Adadelta(params, **optimizer_params)
-
-        # transfer opt state if loaded
-        if opt_name_key in self.loaded_optimizer_states_dict:
-            state = self.loaded_optimizer_states_dict[opt_name_key]
-            optimizer.load_state_dict(state)
-
-        return optimizer

pytorch_lightning/root_module/root_module.py

@@ -5,12 +5,10 @@ import math
 from pytorch_lightning.root_module.memory import ModelSummary
 from pytorch_lightning.root_module.grads import GradInformation
 from pytorch_lightning.root_module.model_saving import ModelIO, load_hparams_from_tags_csv
-from pytorch_lightning.root_module.optimization import OptimizerConfig
 from pytorch_lightning.root_module.hooks import ModelHooks
 
 
-
-class LightningModule(GradInformation, ModelIO, OptimizerConfig, ModelHooks):
+class LightningModule(GradInformation, ModelIO, ModelHooks):
 
     def __init__(self, hparams):
         super(LightningModule, self).__init__()
@@ -21,29 +19,18 @@ class LightningModule(GradInformation, ModelIO, OptimizerConfig, ModelHooks):
         self.current_epoch = 0
         self.global_step = 0
         self.loaded_optimizer_states_dict = {}
-        self.fast_dev_run = hparams.fast_dev_run
-        self.overfit = hparams.overfit
-        self.gradient_clip = hparams.gradient_clip
-        self.num = 2
         self.trainer = None
-        self.from_lightning = True
+        self.experiment = None
+        self.example_input_array = None
 
         # track if gpu was requested for checkpointing
         self.on_gpu = False
-        try:
-            self.on_gpu = hparams.on_gpu
-        except Exception as e:
-            pass
 
         # computed vars for the dataloaders
         self._tng_dataloader = None
         self._val_dataloader = None
         self._test_dataloader = None
 
-        if self.on_gpu:
-            print('running on gpu...')
-            torch.set_default_tensor_type(hparams.default_tensor_type)
-
     def forward(self, *args, **kwargs):
         """
         Expand model in into whatever you need.
@@ -84,15 +71,6 @@ class LightningModule(GradInformation, ModelIO, OptimizerConfig, ModelHooks):
         """
         raise NotImplementedError
 
-    def update_tng_log_metrics(self, logs):
-        """
-        Chance to update metrics to be logged for training step.
-        For example, add music, images, etc... to log
-        :param logs:
-        :return:
-        """
-        raise NotImplementedError
-
     def loss(self, *args, **kwargs):
         """
         Expand model_out into your components
@@ -105,10 +83,6 @@ class LightningModule(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():
             param.requires_grad = False
@@ -141,16 +115,6 @@ class LightningModule(GradInformation, ModelIO, OptimizerConfig, ModelHooks):
         """
         raise NotImplementedError
 
-    @staticmethod
-    def get_process_position(gpus):
-        try:
-            current_gpu = os.environ["CUDA_VISIBLE_DEVICES"]
-            gpu_ids = gpus.split(';')
-            process_position = gpu_ids.index(current_gpu)
-            return process_position, current_gpu
-        except Exception as e:
-            return 0, 0
-
     @classmethod
     def load_from_metrics(cls, weights_path, tags_csv, on_gpu, map_location=None):
         """

pytorch_lightning/testing_models/lm_test_module.py

@@ -1,26 +1,39 @@
+import os
+from collections import OrderedDict
 import torch.nn as nn
-import numpy as np
-from pytorch_lightning.root_module.root_module import LightningModule
-from test_tube import HyperOptArgumentParser
 from torchvision.datasets import MNIST
 import torchvision.transforms as transforms
 import torch
 import torch.nn.functional as F
-import os, pdb
-from collections import OrderedDict
+from test_tube import HyperOptArgumentParser
+from torch import optim
+from torch.utils.data import DataLoader
+from torch.utils.data.distributed import DistributedSampler
+
+from pytorch_lightning.root_module.root_module import LightningModule
 
 
-class ExampleModel(LightningModule):
+class LightningTestModel(LightningModule):
     """
     Sample model to show how to define a template
     """
 
-    def __init__(self, hparams):
+    def __init__(self, hparams, force_remove_distributed_sampler=False):
+        """
+        Pass in parsed HyperOptArgumentParser to the model
+        :param hparams:
+        """
         # init superclass
-        super(ExampleModel, self).__init__(hparams)
+        super(LightningTestModel, self).__init__(hparams)
 
         self.batch_size = hparams.batch_size
 
+        # if you specify an example input, the summary will show input/output for each layer
+        self.example_input_array = torch.rand(5, 28 * 28)
+
+        # remove to test warning for dist sampler
+        self.force_remove_distributed_sampler = force_remove_distributed_sampler
+
         # build model
         self.__build_model()
 
@@ -42,6 +55,11 @@ class ExampleModel(LightningModule):
     # 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)
@@ -59,27 +77,33 @@ class ExampleModel(LightningModule):
 
     def training_step(self, data_batch, batch_i):
         """
-        Called inside the training loop
+        Lightning calls this inside the training loop
         :param data_batch:
         :return:
         """
         # forward pass
         x, y = data_batch
         x = x.view(x.size(0), -1)
+
         y_hat = self.forward(x)
 
         # calculate loss
         loss_val = self.loss(y, y_hat)
 
+        # in DP mode (default) make sure if result is scalar, there's another dim in the beginning
+        if self.trainer.use_dp:
+            loss_val = loss_val.unsqueeze(0)
+
         output = OrderedDict({
-            'loss': loss_val,
-            'tqdm_metrics': {}
+            'loss': loss_val
         })
+
+        # can also return just a scalar instead of a dict (return loss_val)
         return output
 
     def validation_step(self, data_batch, batch_i):
         """
-        Called inside the validation loop
+        Lightning calls this inside the validation loop
         :param data_batch:
         :return:
         """
@@ -92,13 +116,33 @@ class ExampleModel(LightningModule):
         # acc
         labels_hat = torch.argmax(y_hat, dim=1)
         val_acc = torch.sum(y == labels_hat).item() / (len(y) * 1.0)
+        val_acc = torch.tensor(val_acc)
 
-        output = OrderedDict({
-            'val_loss': loss_val,
-            'val_acc': torch.tensor(val_acc),
-        })
-        return output
+        if self.on_gpu:
+            val_acc = val_acc.cuda(loss_val.device.index)
 
+        # in DP mode (default) make sure if result is scalar, there's another dim in the beginning
+        if self.trainer.use_dp:
+            loss_val = loss_val.unsqueeze(0)
+            val_acc = val_acc.unsqueeze(0)
+
+        # alternate possible outputs to test
+        if self.trainer.batch_nb % 1 == 0:
+            output = OrderedDict({
+                'val_loss': loss_val,
+                'val_acc': val_acc,
+            })
+            return output
+        if self.trainer.batch_nb % 2 == 0:
+            return val_acc
+
+        if self.trainer.batch_nb % 3 == 0:
+            output = OrderedDict({
+                'val_loss': loss_val,
+                'val_acc': val_acc,
+                'test_dic': {'val_loss_a': loss_val}
+            })
+            return output
 
     def validation_end(self, outputs):
         """
@@ -106,6 +150,10 @@ class ExampleModel(LightningModule):
         :param outputs: list of individual outputs of each validation step
         :return:
         """
+        # if returned a scalar from validation_step, outputs is a list of tensor scalars
+        # we return just the average in this case (if we want)
+        # return torch.stack(outputs).mean()
+
         val_loss_mean = 0
         val_acc_mean = 0
         for output in outputs:
@@ -114,11 +162,12 @@ class ExampleModel(LightningModule):
 
         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
+    def on_tng_metrics(self, logs):
+        logs['some_tensor_to_test'] = torch.rand(1)
 
     # ---------------------
     # MODEL SAVING
@@ -139,20 +188,31 @@ class ExampleModel(LightningModule):
         return whatever optimizers we want here
         :return: list of optimizers
         """
-        optimizer = self.choose_optimizer(self.hparams.optimizer_name, self.parameters(), {'lr': self.hparams.learning_rate}, 'optimizer')
-        self.optimizers = [optimizer]
-        return self.optimizers
+        optimizer = optim.Adam(self.parameters(), lr=self.hparams.learning_rate)
+        return [optimizer]
 
     def __dataloader(self, train):
         # init data generators
         transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,), (1.0,))])
-
         dataset = MNIST(root=self.hparams.data_root, train=train, transform=transform, download=True)
 
-        loader = torch.utils.data.DataLoader(
+        # when using multi-node we need to add the datasampler
+        train_sampler = None
+        batch_size = self.hparams.batch_size
+
+        try:
+            if self.on_gpu and not self.force_remove_distributed_sampler:
+                train_sampler = DistributedSampler(dataset, rank=self.trainer.proc_rank)
+                batch_size = batch_size // self.trainer.world_size  # scale batch size
+        except Exception as e:
+            pass
+
+        should_shuffle = train_sampler is None
+        loader = DataLoader(
             dataset=dataset,
-            batch_size=self.hparams.batch_size,
-            shuffle=True
+            batch_size=batch_size,
+            shuffle=should_shuffle,
+            sampler=train_sampler
         )
 
         return loader
@@ -189,6 +249,12 @@ class ExampleModel(LightningModule):
 
     @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
@@ -196,16 +262,19 @@ class ExampleModel(LightningModule):
 
         # 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
+        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, type=float, options=[0.0001, 0.0005, 0.001, 0.005],
+        parser.opt_list('--learning_rate', default=0.001*8, type=float, options=[0.0001, 0.0005, 0.001, 0.005],
                         tunable=False)
-        parser.opt_list('--batch_size', default=256, type=int, options=[32, 64, 128, 256], tunable=False)
         parser.opt_list('--optimizer_name', default='adam', type=str, options=['adam'], tunable=False)
+
+        # if using 2 nodes with 4 gpus each the batch size here (256) will be 256 / (2*8) = 16 per gpu
+        parser.opt_list('--batch_size', default=256*8, type=int, options=[32, 64, 128, 256], tunable=False,
+                        help='batch size will be divided over all the gpus being used across all nodes')
         return parser

pytorch_lightning/trainer_main.py

@@ -52,10 +52,6 @@ def main(hparams, cluster, results_dict):
     hparams.__setattr__('nb_gpus', torch.cuda.device_count())
     hparams.__setattr__('inference_mode', hparams.model_load_weights_path is not None)
 
-    # delay each training start to not overwrite logs
-    process_position, current_gpu = TRAINING_MODEL.get_process_position(hparams.gpus)
-    sleep(process_position + 1)
-
     # init experiment
     exp = Experiment(
         name=hparams.tt_name,

pytorch_lightning/utils/arg_parse.py

@@ -39,8 +39,8 @@ def add_default_args(parser, root_dir, rand_seed=None, possible_model_names=None
         parser.add_argument('--model_name', default='', help=','.join(possible_model_names))
 
     # test_tube settings
-    parser.add_argument('-en', '--tt_name', default='r_lib_')
-    parser.add_argument('-td', '--tt_description', default='test research  lib')
+    parser.add_argument('-en', '--tt_name', default='pt_test')
+    parser.add_argument('-td', '--tt_description', default='pytorch lightning test')
     parser.add_argument('--tt_save_path', default=root_dir + '/test_tube_logs', help='logging dir')
     parser.add_argument('--enable_single_run', dest='single_run', action='store_true')
     parser.add_argument('--nb_hopt_trials', default=1, type=int)

pytorch_lightning/utils/debugging.py

@@ -0,0 +1,5 @@
+import pdb
+import sys
+
+class MisconfigurationException(Exception):
+    pass

pytorch_lightning/utils/embeddings.py

@@ -1,104 +0,0 @@
-import torch
-import numpy as np
-from copy import deepcopy
-
-
-class PretrainedEmbedding(torch.nn.Embedding):
-
-    def __init__(self, embedding_path, embedding_dim, task_vocab, freeze=True, *args, **kwargs):
-        """
-        Loads a prebuilt pytorch embedding from any embedding formated file.
-        Padding=0 by default.
-
-        >>> emb = PretrainedEmbedding(embedding_path='glove.840B.300d.txt',embedding_dim=300, task_vocab={'hello': 1, 'world': 2})
-        >>> data = torch.Tensor([[0, 1], [0, 2]]).long()
-        >>> embedded = emb(data)
-        
-
-
-        :param embedding_path:
-        :param emb_dim:
-        :param task_vocab:
-        :param freeze:
-        :return:
-        """
-        # count the vocab
-        self.vocab_size = max(task_vocab.values()) + 1
-        super(PretrainedEmbedding, self).__init__(self.vocab_size, embedding_dim, padding_idx=0, *args, **kwargs)
-
-        # load pretrained embeddings
-        new_emb = self.__load_task_specific_embeddings(deepcopy(task_vocab), embedding_path, embedding_dim, freeze)
-
-        # transfer weights
-        self.weight = new_emb.weight
-
-        # apply freeze
-        should_freeze = not freeze
-        self.weight.requires_grad = should_freeze
-
-    def __load_task_specific_embeddings(self, vocab_words, embedding_path, emb_dim, freeze):
-        """
-        Iterates embedding file to only pull out task specific embeddings
-        :param vocab_words:
-        :param embedding_path:
-        :param emb_dim:
-        :param freeze:
-        :return:
-        """
-
-        # holds final embeddings for relevant words
-        embeddings = np.zeros(shape=(self.vocab_size, emb_dim))
-
-        # load embedding line by line and extract relevant embeddings
-        with open(embedding_path, encoding='utf-8') as f:
-            for line in f:
-                tokens = line.split(' ')
-                word = tokens[0]
-                embedding = tokens[1:]
-                embedding[-1] = embedding[-1][:-1]  # remove last new line
-
-                if word in vocab_words:
-                    vocab_word_i = vocab_words[word]
-
-                    # skip words that try to overwrite pad idx
-                    if vocab_word_i == 0:
-                        del vocab_words[word]
-                        continue
-
-                    emb_vals = np.asarray([float(x) for x in embedding])
-                    embeddings[vocab_word_i] = emb_vals
-
-                    # remove vocab word to early terminate
-                    del vocab_words[word]
-
-                # early break
-                if len(vocab_words) == 0:
-                    break
-
-        # add random vectors for the non-pretrained words
-        # these are vocab words NOT found in the pretrained embeddings
-        for w, i in vocab_words.items():
-            # skip words that try to overwrite pad idx
-            if i == 0:
-                continue
-
-            embedding = np.random.normal(size=emb_dim)
-            embeddings[i] = embedding
-
-        # turn into pt embedding
-        embeddings = torch.FloatTensor(embeddings)
-        embeddings = torch.nn.Embedding.from_pretrained(embeddings, freeze=freeze)
-
-        return embeddings
-
-
-if __name__ == '__main__':
-    emb = PretrainedEmbedding(
-        embedding_path='/Users/waf/Developer',
-        embedding_dim=300,
-        task_vocab={'hello': 1, 'world': 2}
-    )
-
-    data = torch.Tensor([[0, 1], [0, 2]]).long()
-    embedded = emb(data)
-    print(embedded)

pytorch_lightning/utils/plotting.py

@@ -1,28 +0,0 @@
-from matplotlib import pyplot as plt
-import numpy as np
-np.seterr(divide='ignore', invalid='ignore')
-
-
-def plot_confusion_matrix(cm,
-                          save_path,
-                          normalize=False,
-                          title='Confusion matrix',
-                          ylabel='y',
-                          xlabel='x'):
-    """
-    This function prints and plots the confusion matrix.
-    Normalization can be applied by setting `normalize=True`.
-    """
-    if normalize:
-        cm = cm.astype('float') / cm.sum(axis=1)[:, np.newaxis]
-        print("Normalized confusion matrix")
-    else:
-        print('Confusion matrix, without normalization')
-
-    fig = plt.figure()
-    plt.matshow(cm)
-    plt.title(title)
-    plt.colorbar()
-    plt.ylabel(ylabel)
-    plt.xlabel(xlabel)
-    plt.savefig(save_path)

requirements.txt

@@ -1,27 +1,9 @@
-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
-more-itertools==5.0.0
-numpy==1.15.4
-olefile==0.46
-pandas==0.23.4
-Pillow==5.3.0
-pluggy==0.8.0
-py==1.7.0
-pycparser==2.19
-pytest==4.0.2
-python-dateutil==2.7.5
-pytz==2018.7
+coverage==4.5.3
+mkdocs==1.0.4
+pytest==5.0.1
 scikit-learn==0.20.2
-scipy==1.2.0
-six==1.12.0
-sklearn==0.0
-test-tube==0.6282
-torch==1.0.0
-torchvision==0.2.1
-tqdm==4.28.1
+tqdm==4.32.1
+twine==1.13.0
+numpy==1.16.4
+torch>=1.1.0
+torchvision==0.3.0

setup.cfg

@@ -16,6 +16,33 @@ markers =
 ignore = E731,W504
 max-line-length = 120
 
+[coverage:report]
+exclude_lines =
+    pragma: no cover
+    def __repr__
+    if self.debug:
+    if settings.DEBUG
+    raise AssertionError
+    raise NotImplementedError
+    if 0:
+    if __name__ == .__main__.:
+    except Exception as e
+    print(e)
+    print(traceback.print_exc())
+    return *
+    raise Exception
+    warnings
+    print
+    raise RuntimeError
+    break
+    pass
+    os.makedirs
+
+omit =
+    pytorch_lightning/callbacks/pt_callbacks.py
+    tests/test_models.py
+    pytorch_lightning/testing_models/lm_test_module.py
+
 [flake8]
 ignore = E731,W504,F401,F841
 max-line-length = 120

setup.py

@@ -7,7 +7,7 @@ from setuptools import setup, find_packages
 # http://blog.ionelmc.ro/2014/05/25/python-packaging/
 setup(
     name="pytorch-lightning",
-    version='0.111',
+    version='0.3.5',
     description="The Keras for ML researchers using PyTorch",
     author="William Falcon",
     author_email="waf2107@columbia.edu",
@@ -17,9 +17,9 @@ setup(
     keywords=["deep learning", "pytorch", "AI"],
     python_requires=">=3.5",
     install_requires=[
-        "torch>=1.0.0",
+        "torch>=1.1.0",
         "tqdm",
-        "test-tube",
+        "test-tube>=0.6.7.1",
     ],
     packages=find_packages(),
     long_description=open("README.md", encoding="utf-8").read(),

tests/README.md

@@ -0,0 +1,58 @@
+# Pytorch-Lightning Tests
+
+## Running tests
+The automatic travis tests ONLY run CPU-based tests. Although these cover most of the use cases,
+run on a 2-GPU machine to validate the full test-suite.
+
+
+To run all tests do the following:
+```bash
+git clone https://github.com/williamFalcon/pytorch-lightning
+cd pytorch-lightning
+
+# install module locally
+pip install -e .
+
+# install dev deps
+pip install -r requirements.txt
+
+# run tests
+py.test
+```
+
+To test models that require GPU make sure to run the above command on a GPU machine.
+The GPU machine must have:
+1. At least 2 GPUs.
+2. [NVIDIA-apex](https://github.com/NVIDIA/apex#linux) installed.
+
+
+### test_models.py
+This file fits a tiny model on MNIST using these different set-ups.
+1. CPU only.
+2. Single GPU with DP.
+3. Multiple (2) GPUs using DP.
+3. Multiple (2) GPUs using DDP.
+3. Multiple (2) GPUs using DP + apex (for 16-bit precision).
+3. Multiple (2) GPUs using DDP + apex (for 16-bit precision).   
+
+For each set up it also tests:
+1. model saving.
+2. model loading.
+3. predicting with a loaded model.
+4. simulated save from HPC signal.
+5. simulated load from HPC signal.
+
+## Running Coverage   
+
+```bash
+cd pytorch-lightning
+
+# generate coverage 
+pip install coverage
+coverage run tests/test_models.py   
+
+# print coverage stats
+coverage report -m   
+```
+
+

tests/debug.py

@@ -0,0 +1,142 @@
+import pytest
+from pytorch_lightning import Trainer
+from pytorch_lightning.examples.new_project_templates.lightning_module_template import LightningTemplateModel
+from argparse import Namespace
+from test_tube import Experiment
+from pytorch_lightning.callbacks import ModelCheckpoint
+import numpy as np
+import warnings
+import torch
+import os
+import shutil
+import pdb
+
+
+def get_model():
+    # set up model with these hyperparams
+    root_dir = os.path.dirname(os.path.realpath(__file__))
+    hparams = Namespace(**{'drop_prob': 0.2,
+                           'batch_size': 32,
+                           'in_features': 28*28,
+                           'learning_rate': 0.001*8,
+                           'optimizer_name': 'adam',
+                           'data_root': os.path.join(root_dir, 'mnist'),
+                           'out_features': 10,
+                           'hidden_dim': 1000})
+    model = LightningTemplateModel(hparams)
+
+    return model, hparams
+
+
+def get_exp(debug=True):
+    # set up exp object without actually saving logs
+    root_dir = os.path.dirname(os.path.realpath(__file__))
+    exp = Experiment(debug=debug, save_dir=root_dir, name='tests_tt_dir')
+    return exp
+
+
+def init_save_dir():
+    root_dir = os.path.dirname(os.path.realpath(__file__))
+    save_dir = os.path.join(root_dir, 'save_dir')
+
+    if os.path.exists(save_dir):
+        shutil.rmtree(save_dir)
+
+    os.makedirs(save_dir, exist_ok=True)
+
+    return save_dir
+
+
+def clear_save_dir():
+    root_dir = os.path.dirname(os.path.realpath(__file__))
+    save_dir = os.path.join(root_dir, 'save_dir')
+    if os.path.exists(save_dir):
+        shutil.rmtree(save_dir)
+
+
+def load_model(exp, save_dir):
+
+    # load trained model
+    tags_path = exp.get_data_path(exp.name, exp.version)
+    tags_path = os.path.join(tags_path, 'meta_tags.csv')
+
+    checkpoints = [x for x in os.listdir(save_dir) if '.ckpt' in x]
+    weights_dir = os.path.join(save_dir, checkpoints[0])
+
+    trained_model = LightningTemplateModel.load_from_metrics(weights_path=weights_dir, tags_csv=tags_path, on_gpu=True)
+
+    assert trained_model is not None, 'loading model failed'
+
+    return trained_model
+
+
+def run_prediction(dataloader, trained_model):
+    # run prediction on 1 batch
+    for batch in dataloader:
+        break
+
+    x, y = batch
+    x = x.view(x.size(0), -1)
+
+    y_hat = trained_model(x)
+
+    # acc
+    labels_hat = torch.argmax(y_hat, dim=1)
+    val_acc = torch.sum(y == labels_hat).item() / (len(y) * 1.0)
+    val_acc = torch.tensor(val_acc)
+    val_acc = val_acc.item()
+
+    print(val_acc)
+
+    assert val_acc > 0.70, f'this model is expected to get > 0.7 in test set (it got {val_acc})'
+
+
+def mainasdf():
+
+    save_dir = init_save_dir()
+    model, hparams = get_model()
+
+    # exp file to get meta
+    exp = get_exp(False)
+    exp.argparse(hparams)
+    exp.save()
+
+    # exp file to get weights
+    checkpoint = ModelCheckpoint(save_dir)
+
+    trainer = Trainer(
+        experiment=exp,
+        checkpoint_callback=checkpoint,
+        progress_bar=True,
+        max_nb_epochs=1,
+        gpus=[0, 1],
+        distributed_backend='dp',
+        use_amp=True
+    )
+
+    result = trainer.fit(model)
+
+    # correct result and ok accuracy
+    assert result == 1, 'amp + ddp model failed to complete'
+
+    # test model loading
+    pretrained_model = load_model(exp, save_dir)
+
+    # test model preds
+    run_prediction(model.test_dataloader, pretrained_model)
+
+    clear_save_dir()
+
+
+
+
+if __name__ == '__main__':
+    import subprocess
+    import re
+
+    print('getting pid')
+    command = "lsof -i :%s | awk '{print $2}'" % 12910
+    pids = subprocess.check_output(command, shell=True)
+    pids = pids.strip()
+
+    print(len(pids))

tests/test_models.py

@@ -0,0 +1,506 @@
+import pytest
+from pytorch_lightning import Trainer
+from pytorch_lightning.examples.new_project_templates.lightning_module_template import LightningTemplateModel
+from pytorch_lightning.testing_models.lm_test_module import LightningTestModel
+from argparse import Namespace
+from test_tube import Experiment
+from pytorch_lightning.callbacks import ModelCheckpoint, EarlyStopping
+from pytorch_lightning.utils.debugging import MisconfigurationException
+from pytorch_lightning.root_module import memory
+from pytorch_lightning.models.trainer import reduce_distributed_output
+import numpy as np
+import warnings
+import torch
+import os
+import shutil
+
+SEED = 2334
+torch.manual_seed(SEED)
+np.random.seed(SEED)
+
+
+# ------------------------------------------------------------------------
+# TESTS
+# ------------------------------------------------------------------------
+def test_dp_output_reduce():
+
+    # test identity when we have a single gpu
+    out = torch.rand(3, 1)
+    assert reduce_distributed_output(out, nb_gpus=1) is out
+
+    # average when we have multiples
+    assert reduce_distributed_output(out, nb_gpus=2) == out.mean()
+
+    # when we have a dict of vals
+    out = {
+        'a': out,
+        'b': {
+            'c': out
+        }
+    }
+    reduced = reduce_distributed_output(out, nb_gpus=3)
+    assert reduced['a'] == out['a']
+    assert reduced['b']['c'] == out['b']['c']
+
+
+def test_amp_gpu_ddp_slurm_managed():
+    """
+    Make sure DDP + AMP work
+    :return:
+    """
+    if not torch.cuda.is_available():
+        warnings.warn('test_amp_gpu_ddp cannot run. Rerun on a GPU node to run this test')
+        return
+    if not torch.cuda.device_count() > 1:
+        warnings.warn('test_amp_gpu_ddp cannot run. Rerun on a node with 2+ GPUs to run this test')
+        return
+
+    os.environ['MASTER_PORT'] = str(np.random.randint(12000, 19000, 1)[0])
+
+    hparams = get_hparams()
+    model = LightningTestModel(hparams)
+
+    trainer_options = dict(
+        progress_bar=True,
+        max_nb_epochs=1,
+        gpus=[0],
+        distributed_backend='ddp',
+        use_amp=True
+    )
+
+    save_dir = init_save_dir()
+
+    # exp file to get meta
+    exp = get_exp(False)
+    exp.argparse(hparams)
+    exp.save()
+
+    # exp file to get weights
+    checkpoint = ModelCheckpoint(save_dir)
+
+    # add these to the trainer options
+    trainer_options['checkpoint_callback'] = checkpoint
+    trainer_options['experiment'] = exp
+
+    # fit model
+    trainer = Trainer(**trainer_options)
+    trainer.is_slurm_managing_tasks = True
+    result = trainer.fit(model)
+
+    # correct result and ok accuracy
+    assert result == 1, 'amp + ddp model failed to complete'
+
+    # test root model address
+    assert trainer.resolve_root_node_address('abc') == 'abc'
+    assert trainer.resolve_root_node_address('abc[23]') == 'abc23'
+    assert trainer.resolve_root_node_address('abc[23-24]') == 'abc23'
+    assert trainer.resolve_root_node_address('abc[23-24, 45-40, 40]') == 'abc23'
+
+    # test model loading with a map_location
+    map_location = 'cuda:1'
+    pretrained_model = load_model(exp, save_dir, True, map_location)
+
+    # test model preds
+    run_prediction(model.test_dataloader, pretrained_model)
+
+    if trainer.use_ddp:
+        # on hpc this would work fine... but need to hack it for the purpose of the test
+        trainer.model = pretrained_model
+        trainer.optimizers = pretrained_model.configure_optimizers()
+
+    # test HPC loading / saving
+    trainer.hpc_save(save_dir, exp)
+    trainer.hpc_load(save_dir, on_gpu=True)
+
+    # test freeze on gpu
+    model.freeze()
+    model.unfreeze()
+
+    clear_save_dir()
+
+
+def test_early_stopping_cpu_model():
+    """
+    Test each of the trainer options
+    :return:
+    """
+
+    stopping = EarlyStopping()
+    trainer_options = dict(
+        early_stop_callback=stopping,
+        gradient_clip=1.0,
+        overfit_pct=0.20,
+        track_grad_norm=2,
+        print_nan_grads=True,
+        progress_bar=False,
+        experiment=get_exp(),
+        train_percent_check=0.1,
+        val_percent_check=0.1
+    )
+
+    model, hparams = get_model()
+    run_gpu_model_test(trainer_options, model, hparams, on_gpu=False)
+
+    # test freeze on cpu
+    model.freeze()
+    model.unfreeze()
+
+
+def test_cpu_model_with_amp():
+    """
+    Make sure model trains on CPU
+    :return:
+    """
+
+    trainer_options = dict(
+        progress_bar=False,
+        experiment=get_exp(),
+        max_nb_epochs=1,
+        train_percent_check=0.4,
+        val_percent_check=0.4,
+        use_amp=True
+    )
+
+    model, hparams = get_model()
+
+    with pytest.raises((MisconfigurationException, ModuleNotFoundError)):
+        run_gpu_model_test(trainer_options, model, hparams, on_gpu=False)
+
+
+def test_cpu_model():
+    """
+    Make sure model trains on CPU
+    :return:
+    """
+
+    trainer_options = dict(
+        progress_bar=False,
+        experiment=get_exp(),
+        max_nb_epochs=1,
+        train_percent_check=0.4,
+        val_percent_check=0.4
+    )
+
+    model, hparams = get_model()
+
+    run_gpu_model_test(trainer_options, model, hparams, on_gpu=False)
+
+
+def test_all_features_cpu_model():
+    """
+    Test each of the trainer options
+    :return:
+    """
+
+    trainer_options = dict(
+        gradient_clip=1.0,
+        overfit_pct=0.20,
+        track_grad_norm=2,
+        print_nan_grads=True,
+        progress_bar=False,
+        experiment=get_exp(),
+        max_nb_epochs=1,
+        train_percent_check=0.4,
+        val_percent_check=0.4
+    )
+
+    model, hparams = get_model()
+    run_gpu_model_test(trainer_options, model, hparams, on_gpu=False)
+
+
+def test_single_gpu_model():
+    """
+    Make sure single GPU works (DP mode)
+    :return:
+    """
+    if not torch.cuda.is_available():
+        warnings.warn('test_single_gpu_model cannot run. Rerun on a GPU node to run this test')
+        return
+    model, hparams = get_model()
+
+    trainer_options = dict(
+        progress_bar=False,
+        max_nb_epochs=1,
+        train_percent_check=0.1,
+        val_percent_check=0.1,
+        gpus=[0]
+    )
+
+    run_gpu_model_test(trainer_options, model, hparams)
+
+
+def test_multi_gpu_model_dp():
+    """
+    Make sure DP works
+    :return:
+    """
+    if not torch.cuda.is_available():
+        warnings.warn('test_multi_gpu_model_dp cannot run. Rerun on a GPU node to run this test')
+        return
+    if not torch.cuda.device_count() > 1:
+        warnings.warn('test_multi_gpu_model_dp cannot run. Rerun on a node with 2+ GPUs to run this test')
+        return
+    model, hparams = get_model()
+    trainer_options = dict(
+        progress_bar=False,
+        max_nb_epochs=1,
+        train_percent_check=0.1,
+        val_percent_check=0.1,
+        gpus='-1'
+    )
+
+    run_gpu_model_test(trainer_options, model, hparams)
+
+    # test memory helper functions
+    memory.get_gpu_memory_map()
+
+
+def test_amp_gpu_dp():
+    """
+    Make sure DP + AMP work
+    :return:
+    """
+    if not torch.cuda.is_available():
+        warnings.warn('test_amp_gpu_dp cannot run. Rerun on a GPU node to run this test')
+        return
+    if not torch.cuda.device_count() > 1:
+        warnings.warn('test_amp_gpu_dp cannot run. Rerun on a node with 2+ GPUs to run this test')
+        return
+    model, hparams = get_model()
+    trainer_options = dict(
+        max_nb_epochs=1,
+        gpus='0, 1',  # test init with gpu string
+        distributed_backend='dp',
+        use_amp=True
+    )
+    with pytest.raises(MisconfigurationException):
+        run_gpu_model_test(trainer_options, model, hparams)
+
+
+def test_multi_gpu_model_ddp():
+    """
+    Make sure DDP works
+    :return:
+    """
+    if not torch.cuda.is_available():
+        warnings.warn('test_multi_gpu_model_ddp cannot run. Rerun on a GPU node to run this test')
+        return
+    if not torch.cuda.device_count() > 1:
+        warnings.warn('test_multi_gpu_model_ddp cannot run. Rerun on a node with 2+ GPUs to run this test')
+        return
+
+    os.environ['MASTER_PORT'] = str(np.random.randint(12000, 19000, 1)[0])
+    model, hparams = get_model()
+    trainer_options = dict(
+        progress_bar=False,
+        max_nb_epochs=1,
+        train_percent_check=0.4,
+        val_percent_check=0.2,
+        gpus=[0, 1],
+        distributed_backend='ddp'
+    )
+
+    run_gpu_model_test(trainer_options, model, hparams)
+
+
+def test_amp_gpu_ddp():
+    """
+    Make sure DDP + AMP work
+    :return:
+    """
+    if not torch.cuda.is_available():
+        warnings.warn('test_amp_gpu_ddp cannot run. Rerun on a GPU node to run this test')
+        return
+    if not torch.cuda.device_count() > 1:
+        warnings.warn('test_amp_gpu_ddp cannot run. Rerun on a node with 2+ GPUs to run this test')
+        return
+
+    os.environ['MASTER_PORT'] = str(np.random.randint(12000, 19000, 1)[0])
+
+    hparams = get_hparams()
+    model = LightningTestModel(hparams)
+
+    trainer_options = dict(
+        progress_bar=True,
+        max_nb_epochs=1,
+        gpus=[0, 1],
+        distributed_backend='ddp',
+        use_amp=True
+    )
+
+    run_gpu_model_test(trainer_options, model, hparams)
+
+
+def test_ddp_sampler_error():
+    """
+    Make sure DDP + AMP work
+    :return:
+    """
+    if not torch.cuda.is_available():
+        warnings.warn('test_amp_gpu_ddp cannot run. Rerun on a GPU node to run this test')
+        return
+    if not torch.cuda.device_count() > 1:
+        warnings.warn('test_amp_gpu_ddp cannot run. Rerun on a node with 2+ GPUs to run this test')
+        return
+
+    os.environ['MASTER_PORT'] = str(np.random.randint(12000, 19000, 1)[0])
+
+    hparams = get_hparams()
+    model = LightningTestModel(hparams, force_remove_distributed_sampler=True)
+
+    exp = get_exp(True)
+    exp.save()
+
+    trainer = Trainer(
+        experiment=exp,
+        progress_bar=False,
+        max_nb_epochs=1,
+        gpus=[0, 1],
+        distributed_backend='ddp',
+        use_amp=True
+    )
+
+    with pytest.raises(MisconfigurationException):
+        trainer.get_dataloaders(model)
+
+    clear_save_dir()
+
+
+# ------------------------------------------------------------------------
+# UTILS
+# ------------------------------------------------------------------------
+def run_gpu_model_test(trainer_options, model, hparams, on_gpu=True):
+    save_dir = init_save_dir()
+
+    # exp file to get meta
+    exp = get_exp(False)
+    exp.argparse(hparams)
+    exp.save()
+
+    # exp file to get weights
+    checkpoint = ModelCheckpoint(save_dir)
+
+    # add these to the trainer options
+    trainer_options['checkpoint_callback'] = checkpoint
+    trainer_options['experiment'] = exp
+
+    # fit model
+    trainer = Trainer(**trainer_options)
+    result = trainer.fit(model)
+
+    # correct result and ok accuracy
+    assert result == 1, 'amp + ddp model failed to complete'
+
+    # test model loading
+    pretrained_model = load_model(exp, save_dir, on_gpu)
+
+    # test model preds
+    run_prediction(model.test_dataloader, pretrained_model)
+
+    if trainer.use_ddp:
+        # on hpc this would work fine... but need to hack it for the purpose of the test
+        trainer.model = pretrained_model
+        trainer.optimizers = pretrained_model.configure_optimizers()
+
+    # test HPC loading / saving
+    trainer.hpc_save(save_dir, exp)
+    trainer.hpc_load(save_dir, on_gpu=on_gpu)
+
+    clear_save_dir()
+
+
+def get_hparams():
+    root_dir = os.path.dirname(os.path.realpath(__file__))
+    hparams = Namespace(**{'drop_prob': 0.2,
+                           'batch_size': 32,
+                           'in_features': 28*28,
+                           'learning_rate': 0.001*8,
+                           'optimizer_name': 'adam',
+                           'data_root': os.path.join(root_dir, 'mnist'),
+                           'out_features': 10,
+                           'hidden_dim': 1000})
+    return hparams
+
+
+def get_model():
+    # set up model with these hyperparams
+    hparams = get_hparams()
+    model = LightningTemplateModel(hparams)
+
+    return model, hparams
+
+
+def get_exp(debug=True):
+    # set up exp object without actually saving logs
+    root_dir = os.path.dirname(os.path.realpath(__file__))
+    exp = Experiment(debug=debug, save_dir=root_dir, name='tests_tt_dir')
+    return exp
+
+
+def init_save_dir():
+    root_dir = os.path.dirname(os.path.realpath(__file__))
+    save_dir = os.path.join(root_dir, 'save_dir')
+
+    if os.path.exists(save_dir):
+        shutil.rmtree(save_dir)
+
+    os.makedirs(save_dir, exist_ok=True)
+
+    return save_dir
+
+
+def clear_save_dir():
+    root_dir = os.path.dirname(os.path.realpath(__file__))
+    save_dir = os.path.join(root_dir, 'save_dir')
+    if os.path.exists(save_dir):
+        shutil.rmtree(save_dir)
+
+
+def load_model(exp, save_dir, on_gpu, map_location=None):
+
+    # load trained model
+    tags_path = exp.get_data_path(exp.name, exp.version)
+    tags_path = os.path.join(tags_path, 'meta_tags.csv')
+
+    checkpoints = [x for x in os.listdir(save_dir) if '.ckpt' in x]
+    weights_dir = os.path.join(save_dir, checkpoints[0])
+
+    trained_model = LightningTemplateModel.load_from_metrics(weights_path=weights_dir,
+                                                             tags_csv=tags_path,
+                                                             on_gpu=on_gpu,
+                                                             map_location=map_location)
+
+    assert trained_model is not None, 'loading model failed'
+
+    return trained_model
+
+
+def run_prediction(dataloader, trained_model):
+    # run prediction on 1 batch
+    for batch in dataloader:
+        break
+
+    x, y = batch
+    x = x.view(x.size(0), -1)
+
+    y_hat = trained_model(x)
+
+    # acc
+    labels_hat = torch.argmax(y_hat, dim=1)
+    val_acc = torch.sum(y == labels_hat).item() / (len(y) * 1.0)
+    val_acc = torch.tensor(val_acc)
+    val_acc = val_acc.item()
+
+    print(val_acc)
+
+    assert val_acc > 0.50, f'this model is expected to get > 0.50 in test set (it got {val_acc})'
+
+
+def assert_ok_acc(trainer):
+    # this model should get 0.80+ acc
+    acc = trainer.tng_tqdm_dic['val_acc']
+    assert acc > 0.50, f'model failed to get expected 0.50 validation accuracy. Got: {acc}'
+
+
+if __name__ == '__main__':
+    pytest.main([__file__])