release v0.3.6.6

Lkhphuc lr sched

README.md

@@ -9,12 +9,10 @@
 <p align="center">
   The Keras for ML researchers using PyTorch. More control. Less boilerplate.    
 </p>
+
 <p align="center">
   <a href="https://badge.fury.io/py/pytorch-lightning"><img src="https://badge.fury.io/py/pytorch-lightning.svg" alt="PyPI version" height="18"></a>
   <a href="https://pepy.tech/project/pytorch-lightning"><img src="https://pepy.tech/badge/pytorch-lightning" alt="PyPI version" height="18"></a>
-</p>   
-
-<p align="center">
   <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>
@@ -37,24 +35,29 @@ When starting a new project the last thing you want to do is recode a training l
 
 With lightning, you guarantee those parts of your code work so you can focus on what the meat of the research: Data and training, validation loop logic. Don't worry about multiple gpus or speeding up your code, lightning will do that for you!
 
+## How do I do use it?   
+
 To use lightning do 2 things:  
 1. [Define a LightningModel](https://williamfalcon.github.io/pytorch-lightning/LightningModule/RequiredTrainerInterface/)         
 ```python
-import pytorch_lightning as ptl
+import os
 import torch
 from torch.nn import functional as F
 from torch.utils.data import DataLoader
 from torchvision.datasets import MNIST
+import torchvision.transforms as transforms
+
+import pytorch_lightning as ptl
 
 class CoolModel(ptl.LightningModule):
 
-    def __init(self):
+    def __init__(self):
         super(CoolModel, self).__init__()
         # not the best model...
         self.l1 = torch.nn.Linear(28 * 28, 10)
 
     def forward(self, x):
-        return torch.relu(self.l1(x))
+        return torch.relu(self.l1(x.view(x.size(0), -1)))
 
     def my_loss(self, y_hat, y):
         return F.cross_entropy(y_hat, y)
@@ -62,7 +65,7 @@ class CoolModel(ptl.LightningModule):
     def training_step(self, batch, batch_nb):
         x, y = batch
         y_hat = self.forward(x)
-        return {'tng_loss': self.my_loss(y_hat, y)}
+        return {'loss': self.my_loss(y_hat, y)}
 
     def validation_step(self, batch, batch_nb):
         x, y = batch
@@ -70,23 +73,23 @@ class CoolModel(ptl.LightningModule):
         return {'val_loss': self.my_loss(y_hat, y)}
 
     def validation_end(self, outputs):
-        avg_loss = torch.stack([x for x in outputs['val_loss']]).mean()
-        return avg_loss
+        avg_loss = torch.stack([x['val_loss'] for x in outputs]).mean()
+        return {'avg_val_loss': avg_loss}
 
     def configure_optimizers(self):
         return [torch.optim.Adam(self.parameters(), lr=0.02)]
 
     @ptl.data_loader
     def tng_dataloader(self):
-        return DataLoader(MNIST('path/to/save', train=True), batch_size=32)
+        return DataLoader(MNIST(os.getcwd(), train=True, download=True, transform=transforms.ToTensor()), batch_size=32)
 
     @ptl.data_loader
     def val_dataloader(self):
-        return DataLoader(MNIST('path/to/save', train=False), batch_size=32)
+        return DataLoader(MNIST(os.getcwd(), train=True, download=True, transform=transforms.ToTensor()), batch_size=32)
 
     @ptl.data_loader
     def test_dataloader(self):
-        return DataLoader(MNIST('path/to/save', train=False), batch_size=32)
+        return DataLoader(MNIST(os.getcwd(), train=True, download=True, transform=transforms.ToTensor()), batch_size=32)
 ```
 
 2. Fit with a [trainer](https://williamfalcon.github.io/pytorch-lightning/Trainer/)    
@@ -95,15 +98,23 @@ from pytorch_lightning import Trainer
 from test_tube import Experiment
 
 model = CoolModel()
+exp = Experiment(save_dir=os.getcwd())
 
-# fit on 32 gpus across 4 nodes
-exp = Experiment(save_dir='some/dir')
-trainer = Trainer(experiment=exp, nb_gpu_nodes=4, gpus=[0,1,2,3,4,5,6,7])
+# train on cpu using only 10% of the data (for demo purposes)
+trainer = Trainer(experiment=exp, max_nb_epochs=1, train_percent_check=0.1)
 
+# train on 4 gpus
+# trainer = Trainer(experiment=exp, max_nb_epochs=1, gpus=[0, 1, 2, 3])
+
+# train on 32 gpus across 4 nodes (make sure to submit appropriate SLURM job)
+# trainer = Trainer(experiment=exp, max_nb_epochs=1, gpus=[0, 1, 2, 3, 4, 5, 6, 7], nb_gpu_nodes=4)
+
+# train (1 epoch only here for demo)
 trainer.fit(model)
 
-# see all experiment metrics here
-# tensorboard --log_dir some/dir   
+# view tensorflow logs 
+print(f'View tensorboard logs by running\ntensorboard --logdir {os.getcwd()}')
+print('and going to http://localhost:6006 on your browser')
 ```
 
 
@@ -220,7 +231,8 @@ And run tensorboard from that dir
 tensorboard --logdir /some/path     
 ```    
 
-## Lightning automatically automates all of the following ([each is also configurable](https://williamfalcon.github.io/pytorch-lightning/Trainer/)):
+## Lightning automates all of the following ([each is also configurable](https://williamfalcon.github.io/pytorch-lightning/Trainer/)):
+
 
 ###### Checkpointing    
 
@@ -254,9 +266,9 @@ tensorboard --logdir /some/path
 ###### Experiment Logging   
 
 - [Display metrics in progress bar](https://williamfalcon.github.io/pytorch-lightning/Trainer/Logging/#display-metrics-in-progress-bar)
-- Log arbitrary metrics
 - [Log metric row every k batches](https://williamfalcon.github.io/pytorch-lightning/Trainer/Logging/#log-metric-row-every-k-batches)
 - [Process position](https://williamfalcon.github.io/pytorch-lightning/Trainer/Logging/#process-position)
+- [Tensorboard support](https://williamfalcon.github.io/pytorch-lightning/Trainer/Logging/#tensorboard-support)
 - [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)
@@ -264,22 +276,25 @@ tensorboard --logdir /some/path
 ###### 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)
+- [Hooks](https://williamfalcon.github.io/pytorch-lightning/Trainer/hooks/)
+- [Learning rate scheduling](https://williamfalcon.github.io/pytorch-lightning/LightningModule/RequiredTrainerInterface/#configure_optimizers)    
+- [Use multiple optimizers (like GANs)](https://williamfalcon.github.io/pytorch-lightning/LightningModule/RequiredTrainerInterface/#configure_optimizers)
 - [Set how much of the training set to check (1-100%)](https://williamfalcon.github.io/pytorch-lightning/Trainer/Training%20Loop/#set-how-much-of-the-training-set-to-check)
 
 ###### Validation loop    
 
 - [Check validation every n epochs](https://williamfalcon.github.io/pytorch-lightning/Trainer/Validation%20loop/#check-validation-every-n-epochs)
+- [Hooks](https://williamfalcon.github.io/pytorch-lightning/Trainer/hooks/)
 - [Set how much of the validation set to check](https://williamfalcon.github.io/pytorch-lightning/Trainer/Validation%20loop/#set-how-much-of-the-validation-set-to-check)
 - [Set how much of the test set to check](https://williamfalcon.github.io/pytorch-lightning/Trainer/Validation%20loop/#set-how-much-of-the-test-set-to-check)
 - [Set validation check frequency within 1 training epoch](https://williamfalcon.github.io/pytorch-lightning/Trainer/Validation%20loop/#set-validation-check-frequency-within-1-training-epoch)
 - [Set the number of validation sanity steps](https://williamfalcon.github.io/pytorch-lightning/Trainer/Validation%20loop/#set-the-number-of-validation-sanity-steps)
 
 
+
 ## Demo
 ```bash
 # install lightning
@@ -287,19 +302,22 @@ pip install pytorch-lightning
 
 # clone lightning for the demo
 git clone https://github.com/williamFalcon/pytorch-lightning.git
-cd examples/new_project_templates/
+cd pytorch_lightning/examples/new_project_templates/
 
-# run demo (on cpu)
-python trainer_gpu_cluster_template.py
+# all of the following demos use the SAME model to show no modification needs to be made to your code
+
+# train on cpu 
+python single_cpu_template.py
+
+# train on multiple-gpus 
+python single_gpu_node_template.py --gpus "0,1"
+
+# train on 32 gpus on a cluster (run on a SLURM managed cluster)
+python multi_node_cluster_template.py --nb_gpu_nodes 4 --gpus '0,1,2,3,4,5,6,7'
 ```
 
-Without changing the model AT ALL, you can run the model on a single gpu, over multiple gpus, or over multiple nodes.
+## Bleeding edge
+If you can't wait for the next release, install the most up to date code with:  
 ```bash
-# run a grid search on two gpus
-python fully_featured_trainer.py --gpus "0;1"
-
-# run single model on multiple gpus
-python fully_featured_trainer.py --gpus "0;1" --interactive
+pip install git+https://github.com/williamFalcon/pytorch-lightning.git@master --upgrade
 ```
-
-

docs/LightningModule/RequiredTrainerInterface.md

@@ -3,7 +3,7 @@
 
 A lightning module is a strict superclass of nn.Module, it provides a standard interface for the trainer to interact with the model.
 
-The easiest thing to do is copy [this template](../../pytorch_lightning/examples/new_project_templates/lightning_module_template.py) and modify accordingly. 
+The easiest thing to do is copy the [minimal example](https://williamfalcon.github.io/pytorch-lightning/LightningModule/RequiredTrainerInterface/#minimal-example) below and modify accordingly. 
 
 Otherwise, to Define a Lightning Module, implement the following methods:
 
@@ -14,8 +14,6 @@ Otherwise, to Define a Lightning Module, implement the following methods:
 - [validation_end](RequiredTrainerInterface.md#validation_end)
     
 - [configure_optimizers](RequiredTrainerInterface.md#configure_optimizers)
-- [get_save_dict](RequiredTrainerInterface.md#get_save_dict)
-- [load_model_specific](RequiredTrainerInterface.md#load_model_specific)
 
 - [tng_dataloader](RequiredTrainerInterface.md#tng_dataloader)
 - [tng_dataloader](RequiredTrainerInterface.md#tng_dataloader)
@@ -23,27 +21,32 @@ Otherwise, to Define a Lightning Module, implement the following methods:
 
 **Optional**:   
 
+- [on_save_checkpoint](RequiredTrainerInterface.md#on_save_checkpoint)
+- [on_load_checkpoint](RequiredTrainerInterface.md#on_load_checkpoint)
 - [update_tng_log_metrics](RequiredTrainerInterface.md#update_tng_log_metrics)
 - [add_model_specific_args](RequiredTrainerInterface.md#add_model_specific_args)
 
 ---
-**Minimal example**
+### Minimal example   
 ```python
-import pytorch_lightning as ptl
+import os
 import torch
 from torch.nn import functional as F
 from torch.utils.data import DataLoader
 from torchvision.datasets import MNIST
+import torchvision.transforms as transforms
+
+import pytorch_lightning as ptl
 
 class CoolModel(ptl.LightningModule):
 
-    def __init(self):
+    def __init__(self):
         super(CoolModel, self).__init__()
         # not the best model...
         self.l1 = torch.nn.Linear(28 * 28, 10)
 
     def forward(self, x):
-        return torch.relu(self.l1(x))
+        return torch.relu(self.l1(x.view(x.size(0), -1)))
 
     def my_loss(self, y_hat, y):
         return F.cross_entropy(y_hat, y)
@@ -51,7 +54,7 @@ class CoolModel(ptl.LightningModule):
     def training_step(self, batch, batch_nb):
         x, y = batch
         y_hat = self.forward(x)
-        return {'tng_loss': self.my_loss(y_hat, y)}
+        return {'loss': self.my_loss(y_hat, y)}
 
     def validation_step(self, batch, batch_nb):
         x, y = batch
@@ -59,23 +62,23 @@ class CoolModel(ptl.LightningModule):
         return {'val_loss': self.my_loss(y_hat, y)}
 
     def validation_end(self, outputs):
-        avg_loss = torch.stack([x for x in outputs['val_loss']]).mean()
-        return avg_loss
+        avg_loss = torch.stack([x['val_loss'] for x in outputs]).mean()
+        return {'avg_val_loss': avg_loss}
 
     def configure_optimizers(self):
         return [torch.optim.Adam(self.parameters(), lr=0.02)]
 
     @ptl.data_loader
     def tng_dataloader(self):
-        return DataLoader(MNIST('path/to/save', train=True), batch_size=32)
+        return DataLoader(MNIST(os.getcwd(), train=True, download=True, transform=transforms.ToTensor()), batch_size=32)
 
     @ptl.data_loader
     def val_dataloader(self):
-        return DataLoader(MNIST('path/to/save', train=False), batch_size=32)
+        return DataLoader(MNIST(os.getcwd(), train=True, download=True, transform=transforms.ToTensor()), batch_size=32)
 
     @ptl.data_loader
     def test_dataloader(self):
-        return DataLoader(MNIST('path/to/save', train=False), batch_size=32)
+        return DataLoader(MNIST(os.getcwd(), train=True, download=True, transform=transforms.ToTensor()), batch_size=32)
 ```
 
 ---
@@ -222,57 +225,59 @@ def validation_end(self, outputs):
 def configure_optimizers(self)
 ```
 
-Set up as many optimizers as you need. Normally you'd need one. But in the case of GANs or something more esoteric you might have multiple. 
-Lightning will call .backward() and .step() on each one.  If you use 16 bit precision it will also handle that.
+Set up as many optimizers and (optionally) learning rate schedulers as you need. Normally you'd need one. But in the case of GANs or something more esoteric you might have multiple. 
+Lightning will call .backward() and .step() on each one in every epoch.  If you use 16 bit precision it will also handle that.
 
 
 ##### Return
-List - List of optimizers
+List or Tuple - List of optimizers with an optional second list of learning-rate schedulers
 
 **Example**
 
 ``` {.python}
 # most cases
 def configure_optimizers(self):
-    opt = Adam(lr=0.01)
+    opt = Adam(self.parameters(), lr=0.01)
     return [opt]
     
-# gan example
+# gan example, with scheduler for discriminator
 def configure_optimizers(self):
-    generator_opt = Adam(lr=0.01)
-    disriminator_opt = Adam(lr=0.02)
-    return [generator_opt, disriminator_opt] 
+    generator_opt = Adam(self.model_gen.parameters(), lr=0.01)
+    disriminator_opt = Adam(self.model_disc.parameters(), lr=0.02)
+    discriminator_sched = CosineAnnealing(discriminator_opt, T_max=10)
+    return [generator_opt, disriminator_opt], [discriminator_sched]
 ```
 
 --- 
-### get_save_dict 
+### on_save_checkpoint 
 
 ``` {.python}
-def get_save_dict(self)
+def on_save_checkpoint(self, checkpoint)
 ```
-Called by lightning to checkpoint your model. Lightning saves current epoch, current batch nb, etc...
-All you have to return is what specifically about your lightning model you want to checkpoint.
+Called by lightning to checkpoint your model. Lightning saves the training state (current epoch, global_step, etc)
+and also saves the model state_dict. If you want to save anything else, use this method to add your own
+key-value pair.
 
 ##### Return
-Dictionary - No required keys. Most of the time as described in this example.   
+Nothing
 
 **Example**
 
 ``` {.python}
-def get_save_dict(self):
-    # 99% of use cases this is all you need to return
-    checkpoint = {'state_dict': self.state_dict()}
-    return checkpoint
+def on_save_checkpoint(self, checkpoint):
+    # 99% of use cases you don't need to implement this method 
+    checkpoint['something_cool_i_want_to_save'] = my_cool_pickable_object
 ```
 
 --- 
-### load_model_specific 
+### on_load_checkpoint 
 
 ``` {.python}
-def load_model_specific(self, checkpoint)
+def on_load_checkpoint(self, checkpoint)
 ```
-Called by lightning to restore your model. This is your chance to restore your model using the keys you added in get_save_dict.
-Lightning will automatically restore current epoch, batch nb, etc. 
+Called by lightning to restore your model. Lighting auto-restores global step, epoch, etc...
+It also restores the model state_dict.
+If you saved something with **on_save_checkpoint** this is your chance to restore this.
 
 ##### Return
 Nothing  
@@ -280,9 +285,9 @@ Nothing
 **Example**
 
 ``` {.python}
-def load_model_specific(self, checkpoint):
-    # you defined 'state_dict' in get_save_dict()
-    self.load_state_dict(checkpoint['state_dict'])
+def on_load_checkpoint(self, checkpoint):
+    # 99% of the time you don't need to implement this method
+    self.something_cool_i_want_to_save = checkpoint['something_cool_i_want_to_save']
 ```
 
 --- 
@@ -427,4 +432,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

@@ -21,7 +21,8 @@ pretrained_model = MyLightningModule.load_from_metrics(
     map_location=None
 )
     
-# predict 
+# predict
+pretrained_model.eval()
 pretrained_model.freeze()
 y_hat = pretrained_model(x)
 ```

docs/Trainer/Distributed training.md

@@ -23,6 +23,16 @@ 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).   
 
+---
+#### CUDA flags   
+CUDA flags make certain GPUs visible to your script. 
+Lightning sets these for you automatically, there's NO NEED to do this yourself.
+```python
+# lightning will set according to what you give the trainer
+# os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
+# os.environ["CUDA_VISIBLE_DEVICES"] = "0"
+```
+
 ---
 #### 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.    
@@ -43,10 +53,6 @@ trainer = Trainer(amp_level='O2', use_amp=False)
 #### Single-gpu
 Make sure you're on a GPU machine. 
 ```python
-# set these flags
-os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
-os.environ["CUDA_VISIBLE_DEVICES"] = "0"
-
 # DEFAULT
 trainer = Trainer(gpus=[0])
 ```
@@ -56,13 +62,6 @@ trainer = Trainer(gpus=[0])
 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
-# 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"
-
-
 # to use DataParallel (default)
 trainer = Trainer(gpus=[0,1,2,3,4,5,6,7], distributed_backend='dp')
 

docs/Trainer/Logging.md

@@ -50,6 +50,33 @@ exp = Experiment(create_git_tag=True)
 Trainer(experiment=exp)
 ```
 
+---
+### Tensorboard support   
+The experiment object is a strict subclass of Pytorch SummaryWriter. However, this class
+also snapshots every detail about the experiment (data folder paths, code, hyperparams),
+and allows you to visualize it using tensorboard.
+``` {.python}
+from test_tube import Experiment, HyperOptArgumentParser
+
+# exp hyperparams
+args = HyperOptArgumentParser()
+hparams = args.parse_args()
+
+# this is a summaryWriter with nicer logging structure
+exp = Experiment(save_dir='/some/path', create_git_tag=True)
+
+# track experiment details (must be ArgumentParser or HyperOptArgumentParser).
+# each option in the parser is tracked
+exp.argparse(hparams)
+exp.tag({'description': 'running demo'})
+
+# trainer uses the exp object to log exp data
+trainer = Trainer(experiment=exp)
+trainer.fit(model)
+
+# view logs at:
+# tensorboard --logdir /some/path   
+```
 
 ---
 #### Write logs file to csv every k batches 

docs/Trainer/Training Loop.md

@@ -1,4 +1,4 @@
-The lightning training loop handles everything except the actual computations of your model. To decide what will happen in your training loop, define the [training_step function](../../Pytorch-lightning/LightningModule/#training_step).
+The lightning training loop handles everything except the actual computations of your model. To decide what will happen in your training loop, define the [training_step function](https://williamfalcon.github.io/pytorch-lightning/LightningModule/RequiredTrainerInterface/#training_step).
 
 Below are all the things lightning automates for you in the training loop.
 
@@ -11,17 +11,6 @@ Accumulated gradients runs K small batches of size N before doing a backwards pa
 trainer = Trainer(accumulate_grad_batches=1)
 ```
 
----
-#### Anneal Learning rate
-Cut the learning rate by 10 at every epoch listed in this list.
-``` {.python}
-# DEFAULT (don't anneal)
-trainer = Trainer(lr_scheduler_milestones=None)
-
-# cut LR by 10 at 100, 200, and 300 epochs 
-trainer = Trainer(lr_scheduler_milestones='100, 200, 300')
-```
-
 ---
 #### Force training for min or max epochs
 It can be useful to force training for a minimum number of epochs or limit to a max number

docs/Trainer/Validation loop.md

@@ -1,4 +1,4 @@
-The lightning validation loop handles everything except the actual computations of your model. To decide what will happen in your validation loop, define the [validation_step function](../../Pytorch-lightning/LightningModule/#validation_step).
+The lightning validation loop handles everything except the actual computations of your model. To decide what will happen in your validation loop, define the [validation_step function](https://williamfalcon.github.io/pytorch-lightning/LightningModule/RequiredTrainerInterface/#validation_step).
 Below are all the things lightning automates for you in the validation loop.
 
 **Note**   

docs/Trainer/hooks.md

@@ -0,0 +1,86 @@
+# Hooks
+[[Github Code](https://github.com/williamFalcon/pytorch-lightning/blob/master/pytorch_lightning/root_module/hooks.py)]   
+
+There are cases when you might want to do something different at different parts of the training/validation loop.
+To enable a hook, simply override the method in your LightningModule and the trainer will call it at the correct time.
+
+**Contributing** If there's a hook you'd like to add, simply:    
+1. Fork PytorchLightning.    
+2. Add the hook [here](https://github.com/williamFalcon/pytorch-lightning/blob/master/pytorch_lightning/root_module/hooks.py).       
+3. Add the correct place in the [Trainer](https://github.com/williamFalcon/pytorch-lightning/blob/master/pytorch_lightning/models/trainer.py) where it should be called.    
+
+---
+#### on_epoch_start
+Called in the training loop at the very beginning of the epoch.   
+```python
+def on_epoch_start(self):
+    # do something when the epoch starts
+```
+
+---
+#### on_batch_end
+Called in the training loop at the very end of the epoch.   
+```python
+def on_epoch_end(self):
+    # do something when the epoch ends 
+```
+
+---
+#### on_batch_start
+Called in the training loop before anything happens for that batch.   
+```python
+def on_batch_start(self):
+    # do something when the batch starts
+```
+
+---
+#### on_pre_performance_check
+Called at the very beginning of the validation loop.   
+```python
+def on_pre_performance_check(self):
+    # do something before validation starts 
+```
+
+---
+#### on_post_performance_check
+Called at the very end of the validation loop.   
+```python
+def on_post_performance_check(self):
+    # do something before validation end
+```
+
+---
+#### on_tng_metrics
+Called in the training loop, right before metrics are logged.
+Although you can log at any time by using self.experiment, you can use
+this callback to modify what will be logged.
+```python
+def on_tng_metrics(self, metrics):
+    # do something before validation end
+```
+
+---
+#### on_before_zero_grad
+Called in the training loop after taking an optimizer step and before zeroing grads.
+Good place to inspect weight information with weights updated.
+
+Called once per optimizer
+```python
+def on_before_zero_grad(self, optimizer):
+    # do something with the optimizer or inspect it. 
+```
+
+---
+#### on_after_backward
+Called in the training loop after model.backward()
+This is the ideal place to inspect or log gradient information 
+```python
+def on_after_backward(self):
+    # example to inspect gradient information in tensorboard
+    if self.trainer.global_step % 25 == 0:  # don't make the tf file huge
+        params = self.state_dict()
+        for k, v in params.items():
+            grads = v
+            name = k
+            self.experiment.add_histogram(tag=name, values=grads, global_step=self.trainer.global_step)
+```

docs/Trainer/index.md

@@ -49,25 +49,28 @@ But of course the fun is in all the advanced things it can do:
 **Experiment Logging**   
 
 - [Display metrics in progress bar](Logging/#display-metrics-in-progress-bar)
-- Log arbitrary metrics
 - [Log metric row every k batches](Logging/#log-metric-row-every-k-batches)
 - [Process position](Logging/#process-position)
+- [Tensorboard support](https://williamfalcon.github.io/pytorch-lightning/Trainer/Logging/#tensorboard-support)
 - [Save a snapshot of all hyperparameters](Logging/#save-a-snapshot-of-all-hyperparameters) 
 - [Snapshot code for a training run](Logging/#snapshot-code-for-a-training-run) 
 - [Write logs file to csv every k batches](Logging/#write-logs-file-to-csv-every-k-batches)
 
 **Training loop**    
 
-- [Accumulate gradients](Training%20Loop/#accumulated-gradients)
-- [Anneal Learning rate](Training%20Loop/#anneal-learning-rate)
-- [Force training for min or max epochs](Training%20Loop/#force-training-for-min-or-max-epochs)
-- [Force disable early stop](Training%20Loop/#force-disable-early-stop)
-- [Use multiple optimizers (like GANs)](../Pytorch-lightning/LightningModule/#configure_optimizers)
-- [Set how much of the training set to check (1-100%)](Training%20Loop/#set-how-much-of-the-training-set-to-check)
+- [Accumulate gradients](https://williamfalcon.github.io/pytorch-lightning/Trainer/Training%20Loop/#accumulated-gradients)
+- [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)
+- [Hooks](hooks)
+- [Learning rate scheduling](https://williamfalcon.github.io/pytorch-lightning/LightningModule/RequiredTrainerInterface/#configure_optimizers)
+- [Use multiple optimizers (like GANs)](https://williamfalcon.github.io/pytorch-lightning/Pytorch-Lightning/LightningModule/#configure_optimizers)
+- [Set how much of the training set to check (1-100%)](https://williamfalcon.github.io/pytorch-lightning/Trainer/Training%20Loop/#set-how-much-of-the-training-set-to-check)
 
 **Validation loop**    
 
 - [Check validation every n epochs](Validation%20Loop/#check-validation-every-n-epochs)
+- [Hooks](hooks)
 - [Set how much of the validation set to check](Validation%20Loop/#set-how-much-of-the-validation-set-to-check)
 - [Set how much of the test set to check](Validation%20Loop/#set-how-much-of-the-test-set-to-check)
 - [Set validation check frequency within 1 training epoch](Validation%20Loop/#set-validation-check-frequency-within-1-training-epoch)

docs/index.md

@@ -1,10 +1,17 @@
 ###### New project Quick Start    
-To start a new project define these two files.   
+To start a new project you define two files, a LightningModule and a Trainer file.    
 
-1. [Define a LightningModule](/LightningModule/RequiredTrainerInterface/) 
-2. Pick a trainer      
-    - [Basic CPU Trainer](https://github.com/williamFalcon/pytorch-lightning/blob/master/pytorch_lightning/examples/new_project_templates/single_cpu_template.py) 
-    - [GPU cluster Trainer](https://github.com/williamFalcon/pytorch-lightning/blob/master/pytorch_lightning/examples/new_project_templates/multi_node_cluster_template.py)
+A separate trainer file allows to run many LightningModules. Each LightningModule has the core
+logic to a particular research project.    
+
+For example, one lightningModule could be an image classifier, the other
+one could be a seq-2-seq model, both (optionally) ran by the same trainer file.
+
+1. [MNIST LightningModule](https://williamfalcon.github.io/pytorch-lightning/LightningModule/RequiredTrainerInterface/#minimal-example) 
+2. [Trainer](https://williamfalcon.github.io/pytorch-lightning/Trainer/)
+    - [Basic CPU Trainer Template](https://github.com/williamFalcon/pytorch-lightning/blob/master/pytorch_lightning/examples/new_project_templates/single_cpu_template.py) 
+    - [Multi-GPU Trainer Template](https://github.com/williamFalcon/pytorch-lightning/blob/master/pytorch_lightning/examples/new_project_templates/single_gpu_node_template.py)
+    - [GPU cluster Trainer Template](https://github.com/williamFalcon/pytorch-lightning/blob/master/pytorch_lightning/examples/new_project_templates/multi_node_cluster_template.py)
 
 ###### Docs shortcuts
 - [LightningModule](LightningModule/RequiredTrainerInterface/)  
@@ -49,9 +56,9 @@ To start a new project define these two files.
 ###### Experiment Logging   
 
 - [Display metrics in progress bar](https://williamfalcon.github.io/pytorch-lightning/Trainer/Logging/#display-metrics-in-progress-bar)
-- Log arbitrary metrics
 - [Log metric row every k batches](https://williamfalcon.github.io/pytorch-lightning/Trainer/Logging/#log-metric-row-every-k-batches)
 - [Process position](https://williamfalcon.github.io/pytorch-lightning/Trainer/Logging/#process-position)
+- [Tensorboard support](https://williamfalcon.github.io/pytorch-lightning/Trainer/Logging/#tensorboard-support)
 - [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)
@@ -59,16 +66,18 @@ To start a new project define these two files.
 ###### 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)
+- [Hooks](https://williamfalcon.github.io/pytorch-lightning/Trainer/hooks/)
+- [Learning rate scheduling](https://williamfalcon.github.io/pytorch-lightning/LightningModule/RequiredTrainerInterface/#configure_optimizers)
 - [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)
+- [Hooks](https://williamfalcon.github.io/pytorch-lightning/Trainer/hooks/)
 - [Set how much of the validation set to check](https://williamfalcon.github.io/pytorch-lightning/Trainer/Validation%20loop/#set-how-much-of-the-validation-set-to-check)
 - [Set how much of the test set to check](https://williamfalcon.github.io/pytorch-lightning/Trainer/Validation%20loop/#set-how-much-of-the-test-set-to-check)
 - [Set validation check frequency within 1 training epoch](https://williamfalcon.github.io/pytorch-lightning/Trainer/Validation%20loop/#set-validation-check-frequency-within-1-training-epoch)

pytorch_lightning/examples/new_project_templates/lightning_module_template.py

@@ -154,17 +154,6 @@ class LightningTemplateModel(LightningModule):
         tqdm_dic = {'val_loss': val_loss_mean.item(), 'val_acc': val_acc_mean.item()}
         return tqdm_dic
 
-    # ---------------------
-    # MODEL SAVING
-    # ---------------------
-    def get_save_dict(self):
-        checkpoint = {'state_dict': self.state_dict()}
-        return checkpoint
-
-    def load_model_specific(self, checkpoint):
-        self.load_state_dict(checkpoint['state_dict'])
-        pass
-
     # ---------------------
     # TRAINING SETUP
     # ---------------------
@@ -174,7 +163,8 @@ class LightningTemplateModel(LightningModule):
         :return: list of optimizers
         """
         optimizer = optim.Adam(self.parameters(), lr=self.hparams.learning_rate)
-        return [optimizer]
+        scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=10)
+        return [optimizer], [scheduler]
 
     def __dataloader(self, train):
         # init data generators
@@ -231,7 +221,6 @@ class LightningTemplateModel(LightningModule):
         # parser.set_defaults(gradient_clip=5.0)
 
         # network params
-        parser.opt_list('--drop_prob', default=0.2, options=[0.2, 0.5], type=float, tunable=False)
         parser.add_argument('--in_features', default=28*28, type=int)
         parser.add_argument('--out_features', default=10, type=int)
         parser.add_argument('--hidden_dim', default=50000, type=int) # use 500 for CPU, 50000 for GPU to see speed difference

pytorch_lightning/models/sample_model_template/model_template.py

@@ -1,203 +0,0 @@
-import torch.nn as nn
-import numpy as np
-from pytorch_lightning 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
-
-
-class ExampleModel1(LightningModule):
-    """
-    Sample model to show how to define a template
-    """
-
-    def __init__(self, hparams):
-        # init superclass
-        super(ExampleModel1, self).__init__(hparams)
-
-        self.batch_size = hparams.batch_size
-
-        # build model
-        self.__build_model()
-
-    # ---------------------
-    # MODEL SETUP
-    # ---------------------
-    def __build_model(self):
-        """
-        Layout model
-        :return:
-        """
-        self.c_d1 = nn.Linear(in_features=self.hparams.in_features, out_features=self.hparams.hidden_dim)
-        self.c_d1_bn = nn.BatchNorm1d(self.hparams.hidden_dim)
-        self.c_d1_drop = nn.Dropout(self.hparams.drop_prob)
-
-        self.c_d2 = nn.Linear(in_features=self.hparams.hidden_dim, out_features=self.hparams.out_features)
-
-    # ---------------------
-    # TRAINING
-    # ---------------------
-    def forward(self, x):
-        x = self.c_d1(x)
-        x = F.tanh(x)
-        x = self.c_d1_bn(x)
-        x = self.c_d1_drop(x)
-
-        x = self.c_d2(x)
-        logits = F.log_softmax(x, dim=1)
-
-        return logits
-
-    def loss(self, labels, logits):
-        nll = F.nll_loss(logits, labels)
-        return nll
-
-    def training_step(self, data_batch):
-        """
-        Called 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)
-
-        tqdm_dic = {'jefe': 1}
-        return loss_val, tqdm_dic
-
-    def validation_step(self, data_batch):
-        """
-        Called inside the validation loop
-        :param data_batch:
-        :return:
-        """
-        x, y = data_batch
-        x = x.view(x.size(0), -1)
-        y_hat = self.forward(x)
-
-        loss_val = self.loss(y, y_hat)
-
-        # acc
-        labels_hat = torch.argmax(y_hat, dim=1)
-        val_acc = torch.sum(y == labels_hat).item() / (len(y) * 1.0)
-
-        output = {'y_hat': y_hat, 'val_loss': loss_val.item(), 'val_acc': val_acc}
-        return output
-
-    def validation_end(self, outputs):
-        """
-        Called at the end of validation to aggregate outputs
-        :param outputs: list of individual outputs of each validation step
-        :return:
-        """
-        val_loss_mean = 0
-        accs = []
-        for output in outputs:
-            val_loss_mean += output['val_loss']
-            accs.append(output['val_acc'])
-
-        val_loss_mean /= len(outputs)
-        tqdm_dic = {'val_loss': val_loss_mean, 'val_acc': np.mean(accs)}
-        return tqdm_dic
-
-    def update_tng_log_metrics(self, logs):
-        return logs
-
-    # ---------------------
-    # MODEL SAVING
-    # ---------------------
-    def get_save_dict(self):
-        checkpoint = {
-            'state_dict': self.state_dict(),
-        }
-
-        return checkpoint
-
-    def load_model_specific(self, checkpoint):
-        self.load_state_dict(checkpoint['state_dict'])
-        pass
-
-    # ---------------------
-    # TRAINING SETUP
-    # ---------------------
-    def configure_optimizers(self):
-        """
-        return whatever optimizers we want here
-        :return: list of optimizers
-        """
-        optimizer = self.choose_optimizer(self.hparams.optimizer_name, self.parameters(), {'lr': self.hparams.learning_rate}, 'optimizer')
-        self.optimizers = [optimizer]
-        return self.optimizers
-
-    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(
-            dataset=dataset,
-            batch_size=self.hparams.batch_size,
-            shuffle=True
-        )
-
-        return loader
-
-    @data_loader
-    def tng_dataloader(self):
-        if self._tng_dataloader is None:
-            try:
-                self._tng_dataloader = self.__dataloader(train=True)
-            except Exception as e:
-                print(e)
-                raise e
-        return self._tng_dataloader
-
-    @property
-    def val_dataloader(self):
-        if self._val_dataloader is None:
-            try:
-                self._val_dataloader = self.__dataloader(train=False)
-            except Exception as e:
-                print(e)
-                raise e
-        return self._val_dataloader
-
-    @property
-    def test_dataloader(self):
-        if self._test_dataloader is None:
-            try:
-                self._test_dataloader = self.__dataloader(train=False)
-            except Exception as e:
-                print(e)
-                raise e
-        return self._test_dataloader
-
-    @staticmethod
-    def add_model_specific_args(parent_parser):
-        parser = HyperOptArgumentParser(strategy=parent_parser.strategy, parents=[parent_parser])
-
-        # param overwrites
-        # parser.set_defaults(gradient_clip=5.0)
-
-        # network params
-        parser.opt_list('--drop_prob', default=0.2, options=[0.2, 0.5], type=float, tunable=False)
-        parser.add_argument('--in_features', default=28*28)
-        parser.add_argument('--hidden_dim', default=500)
-        parser.add_argument('--out_features', default=10)
-
-        # data
-        parser.add_argument('--data_root', default='/Users/williamfalcon/Developer/personal/research_lib/research_proj/datasets/mnist', type=str)
-
-        # training params (opt)
-        parser.opt_list('--learning_rate', default=0.001, type=float, options=[0.0001, 0.0005, 0.001, 0.005],
-                        tunable=False)
-        parser.opt_list('--batch_size', default=256, type=int, options=[32, 64, 128, 256], tunable=False)
-        parser.opt_list('--optimizer_name', default='adam', type=str, options=['adam'], tunable=False)
-        return parser

pytorch_lightning/models/trainer.py

@@ -10,7 +10,6 @@ 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
@@ -24,7 +23,7 @@ from pytorch_lightning.utils.debugging import MisconfigurationException
 try:
     from apex import amp
     APEX_AVAILABLE = True
-except ModuleNotFoundError: # pragma: no cover
+except Exception:
     APEX_AVAILABLE = False
 
 
@@ -71,7 +70,6 @@ class Trainer(TrainerIO):
                  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,
@@ -104,7 +102,6 @@ class Trainer(TrainerIO):
         :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:
@@ -141,7 +138,6 @@ class Trainer(TrainerIO):
         self.early_stop_callback = early_stop_callback
         self.min_nb_epochs = min_nb_epochs
         self.nb_sanity_val_steps = nb_sanity_val_steps
-        self.lr_scheduler_milestones = [] if lr_scheduler_milestones is None else [int(x.strip()) for x in lr_scheduler_milestones.split(',')]
         self.lr_schedulers = []
         self.amp_level = amp_level
         self.print_nan_grads = print_nan_grads
@@ -161,7 +157,6 @@ class Trainer(TrainerIO):
         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
@@ -443,8 +438,10 @@ class Trainer(TrainerIO):
                 raise MisconfigurationException('amp + cpu is not supported. Please use a GPU option')
 
             # CHOOSE OPTIMIZER
-            # filter out the weights that were done on gpu so we can load on good old cpus
+            # allow for lr schedulers as well
             self.optimizers = model.configure_optimizers()
+            if len(self.optimizers) == 2:
+                self.optimizers, self.lr_schedulers = self.optimizers
 
             self.__run_pretrain_routine(model)
 
@@ -455,8 +452,10 @@ class Trainer(TrainerIO):
     def __dp_train(self, model):
 
         # CHOOSE OPTIMIZER
-        # filter out the weights that were done on gpu so we can load on good old cpus
+        # allow for lr schedulers as well
         self.optimizers = model.configure_optimizers()
+        if len(self.optimizers) == 2:
+            self.optimizers, self.lr_schedulers = self.optimizers
 
         model.cuda(self.data_parallel_device_ids[0])
 
@@ -500,14 +499,19 @@ class Trainer(TrainerIO):
         self.proc_rank = self.node_rank * len(self.data_parallel_device_ids) + gpu_nb
         self.world_size = self.nb_gpu_nodes * len(self.data_parallel_device_ids)
 
+        # let the exp know the rank to avoid overwriting logs
+        self.experiment.rank = self.proc_rank
+
         # set up server using proc 0's ip address
         # try to init for 20 times at max in case ports are taken
         # where to store ip_table
         self.__init_tcp_connection()
 
         # CHOOSE OPTIMIZER
-        # filter out the weights that were done on gpu so we can load on good old cpus
+        # allow for lr schedulers as well
         self.optimizers = model.configure_optimizers()
+        if len(self.optimizers) == 2:
+            self.optimizers, self.lr_schedulers = self.optimizers
 
         # MODEL
         # copy model to each gpu
@@ -587,12 +591,6 @@ class Trainer(TrainerIO):
         # init training constants
         self.__layout_bookeeping()
 
-        # add lr schedulers
-        if self.lr_scheduler_milestones is not None:
-            for optimizer in self.optimizers:
-                scheduler = MultiStepLR(optimizer, self.lr_scheduler_milestones)
-                self.lr_schedulers.append(scheduler)
-
         # print model summary
         if self.proc_rank == 0 and self.print_weights_summary:
             ref_model.summarize()
@@ -608,22 +606,27 @@ class Trainer(TrainerIO):
         if self.proc_rank == 0:
             self.experiment.save()
 
+        # track model now.
+        # if cluster resets state, the model will update with the saved weights
+        self.model = model
+
         # enable cluster checkpointing
+        # also restores training state
         if self.cluster is not None:  # pragma: no cover
             self.enable_auto_hpc_walltime_manager()
 
         # ---------------------------
         # CORE TRAINING LOOP
         # ---------------------------
-        self.model = model
         self.__train()
 
     def __train(self):
         # run all epochs
         for epoch_nb in range(self.current_epoch, self.max_nb_epochs):
             # update the lr scheduler
-            for lr_scheduler in self.lr_schedulers:
-                lr_scheduler.step()
+            if self.lr_schedulers is not None:
+                for lr_scheduler in self.lr_schedulers:
+                    lr_scheduler.step()
 
             model = self.__get_model()
             model.current_epoch = epoch_nb
@@ -769,7 +772,7 @@ class Trainer(TrainerIO):
             output = self.model.training_step(data_batch, batch_nb)
 
         try:
-            model_specific_tqdm_metrics_dic = output['tqdm_metrics']
+            model_specific_tqdm_metrics_dic = output['prog']
         except Exception as e:
             model_specific_tqdm_metrics_dic = {}
 
@@ -854,30 +857,25 @@ class Trainer(TrainerIO):
         elif not can_check_epoch:
             return
 
-        try:
-            # hook
-            if self.__is_function_implemented('on_pre_performance_check'):
-                model = self.__get_model()
-                model.on_pre_performance_check()
+        # hook
+        if self.__is_function_implemented('on_pre_performance_check'):
+            model = self.__get_model()
+            model.on_pre_performance_check()
 
-            # use full val set on end of epoch
-            # use a small portion otherwise
-            max_batches = None if not self.fast_dev_run else 1
-            model_specific_tqdm_metrics_dic = self.validate(
-                self.model,
-                self.val_dataloader,
-                max_batches
-            )
-            self.__add_tqdm_metrics(model_specific_tqdm_metrics_dic)
+        # use full val set on end of epoch
+        # use a small portion otherwise
+        max_batches = None if not self.fast_dev_run else 1
+        model_specific_tqdm_metrics_dic = self.validate(
+            self.model,
+            self.val_dataloader,
+            max_batches
+        )
+        self.__add_tqdm_metrics(model_specific_tqdm_metrics_dic)
 
-            # hook
-            if self.__is_function_implemented('on_post_performance_check'):
-                model = self.__get_model()
-                model.on_post_performance_check()
-
-        except Exception as e:
-            print(e)
-            print(traceback.print_exc())
+        # hook
+        if self.__is_function_implemented('on_post_performance_check'):
+            model = self.__get_model()
+            model.on_post_performance_check()
 
         if self.progress_bar:
             # add model specific metrics
@@ -885,6 +883,6 @@ class Trainer(TrainerIO):
             self.prog_bar.set_postfix(**tqdm_metrics)
 
         # model checkpointing
-        if self.proc_rank == 0 and self.checkpoint_callback:
+        if self.proc_rank == 0 and self.checkpoint_callback is not None:
             print('save callback...')
             self.checkpoint_callback.on_epoch_end(epoch=self.current_epoch, logs=self.__tng_tqdm_dic)

pytorch_lightning/root_module/model_saving.py

@@ -4,34 +4,36 @@ import re
 import pdb
 from pytorch_lightning.pt_overrides.override_data_parallel import LightningDistributedDataParallel, LightningDataParallel
 
+
 class ModelIO(object):
 
-    def load_model_specific(self, checkpoint):
+    def on_load_checkpoint(self, checkpoint):
         """
         Do something with the checkpoint
+        Gives model a chance to load something before state_dict is restored
         :param checkpoint:
         :return:
         """
-        raise NotImplementedError
+        pass
 
-    def get_save_dict(self):
+    def on_save_checkpoint(self, checkpoint):
         """
-        Return specific things for the model
-        :return:
+        Give the model a chance to add something to the checkpoint.
+        state_dict is already there
         """
-        raise NotImplementedError
+        pass
 
     # -------------------------
     # OPTIONAL HOOKS
     # -------------------------
-    def on_hpc_save(self):
+    def on_hpc_save(self, checkpoint):
         """
         Hook to do whatever you need right before Slurm manager saves the model
         :return:
         """
         pass
 
-    def on_hpc_load(self):
+    def on_hpc_load(self, checkpoint):
         """
         Hook to do whatever you need right before Slurm manager loads the model
         :return:
@@ -42,7 +44,6 @@ class ModelIO(object):
 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
@@ -70,24 +71,33 @@ class TrainerIO(object):
             checkpoint['early_stop_callback_wait'] = self.early_stop_callback.wait
             checkpoint['early_stop_callback_patience'] = self.early_stop_callback.patience
 
+        # save optimizers
         optimizer_states = []
         for i, optimizer in enumerate(self.optimizers):
             optimizer_states.append(optimizer.state_dict())
 
         checkpoint['optimizer_states'] = optimizer_states
+        
+        # save lr schedulers
+        lr_schedulers = []
+        for i, scheduler in enumerate(self.lr_schedulers):
+            lr_schedulers.append(scheduler.state_dict())
 
-        # request what to save from the model
+        checkpoint['lr_schedulers'] = lr_schedulers
+
+        # add the state_dict from the model
         model = self.__get_model()
-        checkpoint_dict = model.get_save_dict()
+        checkpoint['state_dict'] = model.state_dict()
+
+        # give the model a chance to add a few things
+        model.on_save_checkpoint(checkpoint)
 
-        # merge trainer and model saving items
-        checkpoint.update(checkpoint_dict)
         return checkpoint
 
     # --------------------
     # HPC IO
     # --------------------
-    def enable_auto_hpc_walltime_manager(self):  # pragma: no cover
+    def enable_auto_hpc_walltime_manager(self):
         if self.cluster is None:
             return
 
@@ -128,6 +138,11 @@ class TrainerIO(object):
         optimizer_states = checkpoint['optimizer_states']
         for optimizer, opt_state in zip(self.optimizers, optimizer_states):
             optimizer.load_state_dict(opt_state)
+        
+        # restore the lr schedulers
+        lr_schedulers = checkpoint['lr_schedulers']
+        for scheduler, lrs_state in zip(self.lr_schedulers, lr_schedulers):
+            scheduler.load_state_dict(lrs_state)
 
     # ----------------------------------
     # PRIVATE OPS
@@ -150,13 +165,14 @@ class TrainerIO(object):
 
         # give model a chance to do something on hpc_save
         model = self.__get_model()
-        model.on_hpc_save()
+        checkpoint = self.dump_checkpoint()
 
-        # request what to save from the model
-        checkpoint_dict = self.dump_checkpoint()
+        model.on_hpc_save(checkpoint)
 
         # do the actual save
-        torch.save(checkpoint_dict, filepath)
+        torch.save(checkpoint, filepath)
+
+        return filepath
 
     def hpc_load(self, folderpath, on_gpu):
         filepath = '{}/hpc_ckpt_{}.ckpt'.format(folderpath, self.max_ckpt_in_folder(folderpath))
@@ -166,15 +182,17 @@ class TrainerIO(object):
         else:
             checkpoint = torch.load(filepath, map_location=lambda storage, loc: storage)
 
-        # load training state
+        # load training state (affects trainer only)
         self.restore_training_state(checkpoint)
 
         # load model state
         model = self.__get_model()
-        model.load_model_specific(checkpoint)
+
+        # load the state_dict on the model automatically
+        model.load_state_dict(checkpoint['state_dict'])
 
         # call model hook
-        model.on_hpc_load()
+        model.on_hpc_load(checkpoint)
 
     def max_ckpt_in_folder(self, path):
         files = os.listdir(path)

pytorch_lightning/root_module/root_module.py

@@ -58,7 +58,7 @@ class LightningModule(GradInformation, ModelIO, ModelHooks):
 
     def configure_optimizers(self):
         """
-        Return array of optimizers
+        Return a list of optimizers and a list of schedulers (could be empty)
         :return:
         """
         raise NotImplementedError
@@ -108,11 +108,13 @@ class LightningModule(GradInformation, ModelIO, ModelHooks):
         else:
             checkpoint = torch.load(weights_path, map_location=lambda storage, loc: storage)
 
+        # load the state_dict on the model automatically
         model = cls(hparams)
+        model.load_state_dict(checkpoint['state_dict'])
+
+        # give model a chance to load something
+        model.on_load_checkpoint(checkpoint)
 
-        # allow model to load
-        model.load_model_specific(checkpoint)
-        model.load_state_dict(checkpoint['state_dict'], strict=False)
         return model
 
     def summarize(self):

pytorch_lightning/testing_models/lm_test_module.py

@@ -96,12 +96,15 @@ class LightningTestModel(LightningModule):
         if self.trainer.use_dp:
             loss_val = loss_val.unsqueeze(0)
 
-        output = OrderedDict({
-            'loss': loss_val
-        })
-
-        # can also return just a scalar instead of a dict (return loss_val)
-        return output
+        # alternate possible outputs to test
+        if self.trainer.batch_nb % 1 == 0:
+            output = OrderedDict({
+                'loss': loss_val,
+                'prog': {'some_val': loss_val * loss_val}
+            })
+            return output
+        if self.trainer.batch_nb % 2 == 0:
+            return loss_val
 
     def validation_step(self, data_batch, batch_i):
         """
@@ -171,17 +174,6 @@ class LightningTestModel(LightningModule):
     def on_tng_metrics(self, logs):
         logs['some_tensor_to_test'] = torch.rand(1)
 
-    # ---------------------
-    # MODEL SAVING
-    # ---------------------
-    def get_save_dict(self):
-        checkpoint = {'state_dict': self.state_dict()}
-        return checkpoint
-
-    def load_model_specific(self, checkpoint):
-        self.load_state_dict(checkpoint['state_dict'])
-        pass
-
     # ---------------------
     # TRAINING SETUP
     # ---------------------
@@ -190,7 +182,10 @@ class LightningTestModel(LightningModule):
         return whatever optimizers we want here
         :return: list of optimizers
         """
+        # try no scheduler for this model (testing purposes)
         optimizer = optim.Adam(self.parameters(), lr=self.hparams.learning_rate)
+
+        # test returning only 1 list instead of 2
         return [optimizer]
 
     def __dataloader(self, train):

pytorch_lightning/trainer_main.py

@@ -1,210 +0,0 @@
-import os
-import sys
-
-import torch
-import numpy as np
-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 time import sleep
-
-from pytorch_lightning.callbacks.pt_callbacks import EarlyStopping, ModelCheckpoint
-SEED = 2334
-torch.manual_seed(SEED)
-np.random.seed(SEED)
-
-# ---------------------
-# DEFINE MODEL HERE
-# ---------------------
-from pytorch_lightning.models.sample_model_template.model_template import ExampleModel1
-# ---------------------
-
-AVAILABLE_MODELS = {
-    'model_1': ExampleModel1
-}
-
-
-"""
-Allows training by using command line arguments
-
-Run by: 
-# TYPE YOUR RUN COMMAND HERE
-"""
-
-
-def main_local(hparams):
-    main(hparams, None, None)
-
-
-def main(hparams, cluster, results_dict):
-    """
-    Main training routine specific for this project
-    :param hparams:
-    :return:
-    """
-    on_gpu = torch.cuda.is_available()
-    if hparams.disable_cuda:
-        on_gpu = False
-
-    device = 'cuda' if on_gpu else 'cpu'
-    hparams.__setattr__('device', device)
-    hparams.__setattr__('on_gpu', on_gpu)
-    hparams.__setattr__('nb_gpus', torch.cuda.device_count())
-    hparams.__setattr__('inference_mode', hparams.model_load_weights_path is not None)
-
-    # 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
-    print('loading model...')
-    model = TRAINING_MODEL(hparams)
-    print('model built')
-
-    # callbacks
-    early_stop = EarlyStopping(
-        monitor=hparams.early_stop_metric,
-        patience=hparams.early_stop_patience,
-        verbose=True,
-        mode=hparams.early_stop_mode
-    )
-
-    model_save_path = '{}/{}/{}'.format(hparams.model_save_path, exp.name, exp.version)
-    checkpoint = ModelCheckpoint(
-        filepath=model_save_path,
-        save_function=None,
-        save_best_only=True,
-        verbose=True,
-        monitor=hparams.model_save_monitor_value,
-        mode=hparams.model_save_monitor_mode
-    )
-
-    # configure trainer
-    trainer = Trainer(
-        experiment=exp,
-        on_gpu=on_gpu,
-        cluster=cluster,
-        progress_bar=hparams.enable_tqdm,
-        overfit_pct=hparams.overfit,
-        track_grad_norm=hparams.track_grad_norm,
-        fast_dev_run=hparams.fast_dev_run,
-        check_val_every_n_epoch=hparams.check_val_every_n_epoch,
-        accumulate_grad_batches=hparams.accumulate_grad_batches,
-        process_position=process_position,
-        current_gpu_name=current_gpu,
-        checkpoint_callback=checkpoint,
-        early_stop_callback=early_stop,
-        enable_early_stop=hparams.enable_early_stop,
-        max_nb_epochs=hparams.max_nb_epochs,
-        min_nb_epochs=hparams.min_nb_epochs,
-        train_percent_check=hparams.train_percent_check,
-        val_percent_check=hparams.val_percent_check,
-        test_percent_check=hparams.test_percent_check,
-        val_check_interval=hparams.val_check_interval,
-        log_save_interval=hparams.log_save_interval,
-        add_log_row_interval=hparams.add_log_row_interval,
-        lr_scheduler_milestones=hparams.lr_scheduler_milestones
-    )
-
-    # 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, 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)
-
-    # use default args
-    root_dir = os.path.split(os.path.dirname(sys.modules['__main__'].__file__))[0]
-    parent_parser = get_default_parser(strategy='random_search', root_dir=root_dir)
-
-    # allow model to overwrite or extend args
-    TRAINING_MODEL = AVAILABLE_MODELS[model_name]
-    parser = TRAINING_MODEL.add_model_specific_args(parent_parser)
-    parser.json_config('-c', '--config', default=root_dir + '/run_configs/local.json')
-    hyperparams = parser.parse_args()
-
-    # format GPU layout
-    os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
-    gpu_ids = hyperparams.gpus.split(';')
-
-    # RUN TRAINING
-    if hyperparams.on_cluster:
-        print('RUNNING ON SLURM CLUSTER')
-        os.environ["CUDA_VISIBLE_DEVICES"] = ','.join(gpu_ids)
-        optimize_on_cluster(hyperparams)
-
-    elif hyperparams.single_run_gpu:
-        print(f'RUNNING 1 TRIAL ON GPU. gpu: {gpu_ids[0]}')
-        os.environ["CUDA_VISIBLE_DEVICES"] = gpu_ids[0]
-        main(hyperparams, None, None)
-
-    elif hyperparams.local or hyperparams.single_run:
-        os.environ["CUDA_VISIBLE_DEVICES"] = '0'
-        print('RUNNING LOCALLY')
-        main(hyperparams, None, None)
-
-    else:
-        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)
-        )

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.3.6',
+    version='0.3.6.6',
     description="The Keras for ML researchers using PyTorch",
     author="William Falcon",
     author_email="waf2107@columbia.edu",
@@ -19,7 +19,7 @@ setup(
     install_requires=[
         "torch>=1.1.0",
         "tqdm",
-        "test-tube>=0.6.7.1",
+        "test-tube>=0.6.7.4",
     ],
     packages=find_packages(),
     long_description=open("README.md", encoding="utf-8").read(),

tests/test_models.py

@@ -3,16 +3,18 @@ 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 test_tube import Experiment, SlurmCluster
 from pytorch_lightning.callbacks import ModelCheckpoint, EarlyStopping
 from pytorch_lightning.utils.debugging import MisconfigurationException
 from pytorch_lightning.root_module import memory
 from pytorch_lightning.models.trainer import reduce_distributed_output
+from pytorch_lightning.root_module import model_saving
 import numpy as np
 import warnings
 import torch
 import os
 import shutil
+import pdb
 
 SEED = 2334
 torch.manual_seed(SEED)
@@ -22,6 +24,135 @@ np.random.seed(SEED)
 # ------------------------------------------------------------------------
 # TESTS
 # ------------------------------------------------------------------------
+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_cpu_slurm_save_load():
+    """
+    Verify model save/load/checkpoint on CPU
+    :return:
+    """
+    hparams = get_hparams()
+    model = LightningTestModel(hparams)
+
+    save_dir = init_save_dir()
+
+    # exp file to get meta
+    exp = get_exp(False)
+    exp.argparse(hparams)
+    exp.save()
+
+    cluster_a = SlurmCluster()
+    trainer_options = dict(
+        max_nb_epochs=1,
+        cluster=cluster_a,
+        experiment=exp,
+        checkpoint_callback=ModelCheckpoint(save_dir)
+    )
+
+    # fit model
+    trainer = Trainer(**trainer_options)
+    result = trainer.fit(model)
+    real_global_step = trainer.global_step
+
+    # traning complete
+    assert result == 1, 'amp + ddp model failed to complete'
+
+    # predict with trained model before saving
+    # make a prediction
+    for batch in model.test_dataloader:
+        break
+
+    x, y = batch
+    x = x.view(x.size(0), -1)
+
+    model.eval()
+    pred_before_saving = model(x)
+
+    # test registering a save function
+    trainer.enable_auto_hpc_walltime_manager()
+
+    # test HPC saving
+    # simulate snapshot on slurm
+    saved_filepath = trainer.hpc_save(save_dir, exp)
+    assert os.path.exists(saved_filepath)
+
+    # wipe-out trainer and model
+    # retrain with not much data... this simulates picking training back up after slurm
+    # we want to see if the weights come back correctly
+    continue_tng_hparams = get_hparams(continue_training=True, hpc_exp_number=cluster_a.hpc_exp_number)
+    trainer_options = dict(
+        max_nb_epochs=1,
+        cluster=SlurmCluster(continue_tng_hparams),
+        experiment=exp,
+        checkpoint_callback=ModelCheckpoint(save_dir),
+    )
+    trainer = Trainer(**trainer_options)
+    model = LightningTestModel(hparams)
+
+    # set the epoch start hook so we can predict before the model does the full training
+    def assert_pred_same():
+        assert trainer.global_step == real_global_step and trainer.global_step > 0
+
+        # predict with loaded model to make sure answers are the same
+        trainer.model.eval()
+        new_pred = trainer.model(x)
+        assert torch.all(torch.eq(pred_before_saving, new_pred)).item() == 1
+
+    model.on_epoch_start = assert_pred_same
+
+    # by calling fit again, we trigger training, loading weights from the cluster
+    # and our hook to predict using current model before any more weight updates
+    trainer.fit(model)
+
+    clear_save_dir()
+
+
+def test_loading_meta_tags():
+    hparams = get_hparams()
+
+    save_dir = init_save_dir()
+
+    # save tags
+    exp = get_exp(False)
+    exp.tag({'some_str':'a_str', 'an_int': 1, 'a_float': 2.0})
+    exp.argparse(hparams)
+    exp.save()
+
+    # load tags
+    tags_path = exp.get_data_path(exp.name, exp.version) + '/meta_tags.csv'
+    tags = model_saving.load_hparams_from_tags_csv(tags_path)
+
+    assert tags.batch_size == 32 and tags.hidden_dim == 1000
+
+    clear_save_dir()
+
+
 def test_dp_output_reduce():
 
     # test identity when we have a single gpu
@@ -43,6 +174,74 @@ def test_dp_output_reduce():
     assert reduced['b']['c'] == out['b']['c']
 
 
+def test_model_saving_loading():
+    """
+    Tests use case where trainer saves the model, and user loads it from tags independently
+    :return:
+    """
+    hparams = get_hparams()
+    model = LightningTestModel(hparams)
+
+    save_dir = init_save_dir()
+
+    # exp file to get meta
+    exp = get_exp(False)
+    exp.argparse(hparams)
+    exp.save()
+
+    trainer_options = dict(
+        max_nb_epochs=1,
+        cluster=SlurmCluster(),
+        experiment=exp,
+        checkpoint_callback=ModelCheckpoint(save_dir)
+    )
+
+    # fit model
+    trainer = Trainer(**trainer_options)
+    result = trainer.fit(model)
+
+    # traning complete
+    assert result == 1, 'amp + ddp model failed to complete'
+
+    # make a prediction
+    for batch in model.test_dataloader:
+        break
+
+    x, y = batch
+    x = x.view(x.size(0), -1)
+
+    # generate preds before saving model
+    model.eval()
+    pred_before_saving = model(x)
+
+    # save model
+    new_weights_path = os.path.join(save_dir, 'save_test.ckpt')
+    trainer.save_checkpoint(new_weights_path)
+
+    # load new model
+    tags_path = exp.get_data_path(exp.name, exp.version)
+    tags_path = os.path.join(tags_path, 'meta_tags.csv')
+    model_2 = LightningTestModel.load_from_metrics(weights_path=new_weights_path, tags_csv=tags_path, on_gpu=False)
+    model_2.eval()
+
+    # make prediction
+    # assert that both predictions are the same
+    new_pred = model_2(x)
+    assert torch.all(torch.eq(pred_before_saving, new_pred)).item() == 1
+
+    clear_save_dir()
+
+
+
+
+def test_model_freeze_unfreeze():
+    hparams = get_hparams()
+    model = LightningTestModel(hparams)
+
+    model.freeze()
+    model.unfreeze()
+
+
 def test_amp_gpu_ddp_slurm_managed():
     """
     Make sure DDP + AMP work
@@ -108,7 +307,7 @@ def test_amp_gpu_ddp_slurm_managed():
     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()
+        trainer.optimizers, trainer.lr_schedulers = pretrained_model.configure_optimizers()
 
     # test HPC loading / saving
     trainer.hpc_save(save_dir, exp)
@@ -305,33 +504,6 @@ def test_multi_gpu_model_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():
     """
@@ -402,7 +574,7 @@ def run_gpu_model_test(trainer_options, model, hparams, on_gpu=True):
     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()
+        trainer.optimizers, trainer.lr_schedulers = pretrained_model.configure_optimizers()
 
     # test HPC loading / saving
     trainer.hpc_save(save_dir, exp)
@@ -411,16 +583,24 @@ def run_gpu_model_test(trainer_options, model, hparams, on_gpu=True):
     clear_save_dir()
 
 
-def get_hparams():
+def get_hparams(continue_training=False, hpc_exp_number=0):
     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})
+
+    args = {
+        'drop_prob': 0.2,
+        'batch_size': 32,
+        'in_features': 28*28,
+        'learning_rate': 0.001*8,
+        'optimizer_name': 'adam',
+        'data_root': os.path.join(root_dir, 'mnist'),
+        'out_features': 10,
+        'hidden_dim': 1000}
+
+    if continue_training:
+        args['test_tube_do_checkpoint_load'] = True
+        args['hpc_exp_number'] = hpc_exp_number
+
+    hparams = Namespace(**args)
     return hparams