release v0.4.6

* added smarter optimizer options

* added smarter optimizer options

* added smarter optimizer options tests

* added smarter optimizer options tests

* added smarter optimizer options tests

* added smarter optimizer options tests

* added smarter optimizer options tests

* added smarter optimizer options tests

* added smarter optimizer options tests

* added single gpu data transfer recursive

* added single gpu data transfer recursive

* added single gpu data transfer recursive

* added single gpu data transfer recursive

* added single gpu data transfer recursive

.github/ISSUE_TEMPLATE/how-to-question.md

@@ -0,0 +1,26 @@
+---
+name: How to question
+about: Asking how-to questions
+title: ''
+labels: question
+assignees: ''
+
+---
+
+### Before asking:   
+1. search the issues.   
+2. search the docs.    
+
+If you still can't find what you need:     
+#### What is your question?    
+
+#### Code    
+Please paste a code snippet if your question requires it!   
+
+#### What have you tried?    
+
+#### What's your environment?   
+- conda version (no venv)   
+- PyTorch version   
+- Lightning version   
+- Test-tube version

.gitignore

@@ -11,6 +11,7 @@ pip-wheel-metadata/
 test_tube_exp/
 tests/tests_tt_dir/
 tests/save_dir
+default/
 
 # Byte-compiled / optimized / DLL files
 __pycache__/

README.md

@@ -10,17 +10,15 @@
 [![PyPI Status](https://badge.fury.io/py/pytorch-lightning.svg)](https://badge.fury.io/py/pytorch-lightning)
 [![PyPI Status](https://pepy.tech/badge/pytorch-lightning)](https://pepy.tech/project/pytorch-lightning)
 [![Build Status](https://travis-ci.org/williamFalcon/pytorch-lightning.svg?branch=master)](https://travis-ci.org/williamFalcon/pytorch-lightning)
-<!-- 
-removed until windows install issues resolved.
-[![Build status](https://ci.appveyor.com/api/projects/status/rum89d7hq8l1kfye?svg=true)](https://ci.appveyor.com/project/Borda/pytorch-lightning) -->
-<!-- 
-removed until codecov badge isn't empy. likely a config error showing nothing on master.
-[![codecov](https://codecov.io/gh/Borda/pytorch-lightning/branch/master/graph/badge.svg)](https://codecov.io/gh/Borda/pytorch-lightning)
--->
+[![Build status](https://ci.appveyor.com/api/projects/status/rum89d7hq8l1kfye?svg=true)](https://ci.appveyor.com/project/Borda/pytorch-lightning)
 [![Coverage](https://github.com/williamFalcon/pytorch-lightning/blob/master/docs/source/_static/coverage.svg)](https://github.com/williamFalcon/pytorch-lightning/tree/master/tests#running-coverage)
 [![CodeFactor](https://www.codefactor.io/repository/github/borda/pytorch-lightning/badge)](https://www.codefactor.io/repository/github/borda/pytorch-lightning)
 [![ReadTheDocs](https://readthedocs.org/projects/pytorch-lightning/badge/?version=latest)](https://pytorch-lightning.readthedocs.io/en/latest)
 [![license](https://img.shields.io/badge/License-Apache%202.0-blue.svg)](https://github.com/williamFalcon/pytorch-lightning/blob/master/LICENSE)
+<!-- 
+removed until codecov badge isn't empy. likely a config error showing nothing on master.
+[![codecov](https://codecov.io/gh/Borda/pytorch-lightning/branch/master/graph/badge.svg)](https://codecov.io/gh/Borda/pytorch-lightning)
+-->
 
 </div>
 
@@ -58,9 +56,12 @@ Don't worry about training on multiple gpus or speeding up your code, lightning
 
 ---   
 ## How do I do use it?   
+The research code goes into a [LightningModule]((https://williamfalcon.github.io/pytorch-lightning/LightningModule/RequiredTrainerInterface/)) which you fit using a Trainer.    
+
+Think of the LightningModule as a *system* such as seq-2-seq, GAN, etc... However, the LightningModule can ALSO just be a simple classifier such as the example below.     
 
 To use lightning do 2 things:  
-1. [Define a LightningModel](https://williamfalcon.github.io/pytorch-lightning/LightningModule/RequiredTrainerInterface/)         
+1. [Define a LightningModule](https://williamfalcon.github.io/pytorch-lightning/LightningModule/RequiredTrainerInterface/)         
 ```python
 import os
 import torch
@@ -71,10 +72,10 @@ import torchvision.transforms as transforms
 
 import pytorch_lightning as pl
 
-class CoolModel(pl.LightningModule):
+class CoolSystem(pl.LightningModule):
 
     def __init__(self):
-        super(CoolModel, self).__init__()
+        super(CoolSystem, self).__init__()
         # not the best model...
         self.l1 = torch.nn.Linear(28 * 28, 10)
 
@@ -91,7 +92,7 @@ class CoolModel(pl.LightningModule):
         # OPTIONAL
         x, y = batch
         y_hat = self.forward(x)
-        return {'val_loss': self.my_loss(y_hat, y)}
+        return {'val_loss': F.cross_entropy(y_hat, y)}
 
     def validation_end(self, outputs):
         # OPTIONAL
@@ -122,7 +123,7 @@ class CoolModel(pl.LightningModule):
 ```python
 from pytorch_lightning import Trainer
 
-model = CoolModel()
+model = CoolSystem()
 
 # most basic trainer, uses good defaults
 trainer = Trainer()    
@@ -137,7 +138,7 @@ from test_tube import Experiment
 exp = Experiment(save_dir=os.getcwd())
 
 # train on cpu using only 10% of the data (for demo purposes)
-# pass in experi
+# pass in experiment for automatic tensorboard logging.    
 trainer = Trainer(experiment=exp, max_nb_epochs=1, train_percent_check=0.1)
 
 # train on 4 gpus
@@ -316,6 +317,7 @@ tensorboard --logdir /some/path
 - [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)
+- [Step optimizers at arbitrary intervals](https://williamfalcon.github.io/pytorch-lightning/Trainer/hooks/#optimizer_step)
 
 ###### Validation loop    
 
@@ -375,8 +377,19 @@ Nope.
 Nope. Please use anaconda or miniconda.    
 
 **Which PyTorch versions do you support?**    
-Lightning 0.4.2+ supports PyTorch 1.2.0.    
-For PyTorch 1.1.0 install Lightning 0.4.0 with test-tube=0.6.7.6.    
+##### PyTorch 1.1.0       
+```bash    
+# install pytorch 1.1.0 using the official instructions   
+
+# install test-tube 0.6.7.6 which supports 1.1.0   
+pip install test-tube==0.6.7.6   
+
+# install latest Lightning version without upgrading deps    
+pip install -U --no-deps pytorch-lightning
+```     
+
+##### PyTorch 1.2.0   
+Install via pip as normal   
 
 ## Bleeding edge
 If you can't wait for the next release, install the most up to date code with:  

appveyor.yml

@@ -45,9 +45,7 @@ install:
   # directly to master instead of just PR builds (or the converse).
   - SET PATH=%PYTHON%;%PYTHON%\\Scripts;%path%
   - pip install -U --user pip
-  - pip install "https://download.pytorch.org/whl/cu90/torch-1.1.0-cp%PIP_PYVER%-cp%PIP_PYVER%m-win_amd%PYTHON_ARCH%.whl"
-    pip install "https://download.pytorch.org/whl/cu90/torchvision-0.3.0-cp%PIP_PYVER%-cp%PIP_PYVER%m-win_amd%PYTHON_ARCH%.whl"
-  - pip install -r requirements.txt
+  - pip install -r requirements.txt -f https://download.pytorch.org/whl/torch_stable.html
   - pip install -r ./tests/requirements.txt
 
 # scripts to run before tests (working directory and environment changes are persisted from the previous steps such as "before_build")

docs/LightningModule/RequiredTrainerInterface.md

@@ -9,20 +9,20 @@ Otherwise, to Define a Lightning Module, implement the following methods:
 
 **Required**:  
 
-- [training_step](RequiredTrainerInterface.md#training_step)   
-- [tng_dataloader](RequiredTrainerInterface.md#tng_dataloader)
-- [configure_optimizers](RequiredTrainerInterface.md#configure_optimizers)
+- [training_step](RequiredTrainerInterface.md#training_step)      
+- [tng_dataloader](RequiredTrainerInterface.md#tng_dataloader)    
+- [configure_optimizers](RequiredTrainerInterface.md#configure_optimizers)    
 
 **Optional**:   
-- [validation_step](RequiredTrainerInterface.md#validation_step)
-- [validation_end](RequiredTrainerInterface.md#validation_end)
-- [val_dataloader](RequiredTrainerInterface.md#val_dataloader)
-- [test_dataloader](RequiredTrainerInterface.md#test_dataloader)
 
-- [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)
+- [validation_step](RequiredTrainerInterface.md#validation_step)    
+- [validation_end](RequiredTrainerInterface.md#validation_end)    
+- [val_dataloader](RequiredTrainerInterface.md#val_dataloader)    
+- [test_dataloader](RequiredTrainerInterface.md#test_dataloader)    
+- [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   
@@ -136,7 +136,17 @@ def training_step(self, data_batch, batch_nb):
     
     # return a dict
     return output
-```
+```    
+
+If you define multiple optimizers, this step will also be called with an additional ```optimizer_idx``` param.    
+``` {.python}
+# Multiple optimizers (ie: GANs)     
+def training_step(self, data_batch, batch_nb, optimizer_idx):
+    if optimizer_idx == 0:
+        # do training_step with encoder
+    if optimizer_idx == 1:
+        # do training_step with decoder    
+```    
 
 --- 
 ### tng_dataloader 
@@ -175,6 +185,9 @@ def configure_optimizers(self)
 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.
 
+**Note:** If you use multiple optimizers, training_step will have an additional ```optimizer_idx``` parameter.    
+
+
 
 ##### Return
 List or Tuple - List of optimizers with an optional second list of learning-rate schedulers
@@ -193,14 +206,20 @@ def configure_optimizers(self):
     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]
-```
+```  
+
+If you need to control how often those optimizers step or override the default .step() schedule, override 
+the [optimizer_step](https://williamfalcon.github.io/pytorch-lightning/Trainer/hooks/#optimizer_step) hook.   
 
 ## Optional Methods    
 
 ### validation_step
 
 ``` {.python}
-def validation_step(self, data_batch, batch_nb, dataloader_i)
+def validation_step(self, data_batch, batch_nb)   
+
+# if have multiple val dataloaders:  
+def validation_step(self, data_batch, batch_nb, dataloader_idx)
 ```
 **OPTIONAL**    
 If you don't need to validate you don't need to implement this method.    
@@ -215,7 +234,7 @@ This is most likely the same as your training_step. But unlike training step, th
 |---|---|
 |  data_batch | The output of your dataloader. A tensor, tuple or list  |
 |  batch_nb | Integer displaying which batch this is  |
-|  dataloader_i | Integer displaying which dataloader this is  |
+|  dataloader_i | Integer displaying which dataloader this is (only if multiple val datasets used)  |
 
 **Return**   
 
@@ -226,6 +245,7 @@ This is most likely the same as your training_step. But unlike training step, th
 **Example**
 
 ``` {.python}
+# CASE 1: A single validation dataset
 def validation_step(self, data_batch, batch_nb):
     x, y, z = data_batch
     
@@ -246,7 +266,17 @@ def validation_step(self, data_batch, batch_nb):
     
     # return an optional dict
     return output
-```
+```   
+
+If you pass in multiple validation datasets, validation_step will have an additional argument.
+
+```python
+# CASE 2: multiple validation datasets
+def validation_step(self, data_batch, batch_nb, dataset_idx):
+    # dataset_idx tells you which dataset this is.   
+```   
+
+The ```dataset_idx``` corresponds to the order of datasets returned in ```val_dataloader```.    
 
 --- 
 ### validation_end
@@ -371,6 +401,9 @@ def val_dataloader(self):
     return [loader_a, loader_b, ..., loader_n]   
 ```
 
+In the case where you return multiple val_dataloaders, the validation_step will have an arguement ```dataset_idx```
+which matches the order here.    
+
 --- 
 ### test_dataloader 
 

docs/Trainer/hooks.md

@@ -67,6 +67,36 @@ def on_tng_metrics(self, metrics):
     # do something before validation end
 ```
 
+---   
+#### optimizer_step 
+Calls .step() and .zero_grad for each optimizer.  
+You can override this method to adjust how you do the optimizer step for each optimizer
+
+Called once per optimizer
+```python
+# DEFAULT
+def optimizer_step(self, current_epoch, batch_nb, optimizer, optimizer_i):
+    optimizer.step()   
+    optimizer.zero_grad()   
+    
+# Alternating schedule for optimizer steps (ie: GANs)    
+def optimizer_step(self, current_epoch, batch_nb, optimizer, optimizer_i):
+    # update generator opt every 2 steps
+    if optimizer_i == 0:
+        if batch_nb % 2 == 0 :
+            optimizer.step()
+            optimizer.zero_grad()
+   
+    # update discriminator opt every 4 steps
+    if optimizer_i == 1:
+        if batch_nb % 4 == 0 :
+            optimizer.step()
+            optimizer.zero_grad()    
+    
+    # ...
+    # add as many optimizers as you want 
+```
+
 ---
 #### on_before_zero_grad
 Called in the training loop after taking an optimizer step and before zeroing grads.

docs/Trainer/index.md

@@ -68,6 +68,7 @@ But of course the fun is in all the advanced things it can do:
 - [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)
+- [Step optimizers at arbitrary intervals](https://williamfalcon.github.io/pytorch-lightning/Trainer/hooks/#optimizer_step)
 
 **Validation loop**    
 

docs/index.md

@@ -75,6 +75,7 @@ one could be a seq-2-seq model, both (optionally) ran by the same trainer file.
 - [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)
+- [Step optimizers at arbitrary intervals](https://williamfalcon.github.io/pytorch-lightning/Trainer/hooks/#optimizer_step)
 
 ###### Validation loop    
 

examples/new_project_templates/lightning_module_template.py

@@ -105,7 +105,7 @@ class LightningTemplateModel(LightningModule):
         # can also return just a scalar instead of a dict (return loss_val)
         return output
 
-    def validation_step(self, data_batch, batch_i, dataloader_i):
+    def validation_step(self, data_batch, batch_i):
         """
         Lightning calls this inside the validation loop
         :param data_batch:
@@ -218,7 +218,7 @@ class LightningTemplateModel(LightningModule):
     @pl.data_loader
     def val_dataloader(self):
         print('val data loader called')
-        return [self.__dataloader(train=False) for i in range(2)]
+        return self.__dataloader(train=False)
 
     @pl.data_loader
     def test_dataloader(self):

examples/templates/gan.py

@@ -0,0 +1,178 @@
+"""
+To run this template just do:  
+python gan.py   
+
+After a few epochs, launch tensorboard to see the images being generated at every batch.   
+
+tensorboard --logdir default
+"""
+from argparse import ArgumentParser
+import os
+import numpy as np
+
+import torchvision
+import torchvision.transforms as transforms
+from torchvision.datasets import MNIST
+
+from torch.utils.data import DataLoader
+
+import torch.nn as nn
+import torch.nn.functional as F
+import torch
+
+import pytorch_lightning as pl
+from test_tube import Experiment
+
+
+class Generator(nn.Module):
+    def __init__(self, latent_dim, img_shape):
+        super(Generator, self).__init__()
+        self.img_shape = img_shape
+
+        def block(in_feat, out_feat, normalize=True):
+            layers = [nn.Linear(in_feat, out_feat)]
+            if normalize:
+                layers.append(nn.BatchNorm1d(out_feat, 0.8))
+            layers.append(nn.LeakyReLU(0.2, inplace=True))
+            return layers
+
+        self.model = nn.Sequential(
+            *block(latent_dim, 128, normalize=False),
+            *block(128, 256),
+            *block(256, 512),
+            *block(512, 1024),
+            nn.Linear(1024, int(np.prod(img_shape))),
+            nn.Tanh()
+        )
+
+    def forward(self, z):
+        img = self.model(z)
+        img = img.view(img.size(0), *self.img_shape)
+        return img
+
+
+class Discriminator(nn.Module):
+    def __init__(self, img_shape):
+        super(Discriminator, self).__init__()
+
+        self.model = nn.Sequential(
+            nn.Linear(int(np.prod(img_shape)), 512),
+            nn.LeakyReLU(0.2, inplace=True),
+            nn.Linear(512, 256),
+            nn.LeakyReLU(0.2, inplace=True),
+            nn.Linear(256, 1),
+            nn.Sigmoid(),
+        )
+
+    def forward(self, img):
+        img_flat = img.view(img.size(0), -1)
+        validity = self.model(img_flat)
+
+        return validity
+
+
+class GAN(pl.LightningModule):
+
+    def __init__(self, hparams):
+        super(GAN, self).__init__()
+        self.hparams = hparams
+
+        # networks
+        mnist_shape = (1, 28, 28)
+        self.generator = Generator(latent_dim=hparams.latent_dim, img_shape=mnist_shape)
+        self.discriminator = Discriminator(img_shape=mnist_shape)
+
+        # cache for generated images
+        self.generated_imgs = None
+
+    def forward(self, z):
+        return self.generator(z)
+
+    def adversarial_loss(self, y_hat, y):
+        return F.binary_cross_entropy(y_hat, y)
+
+    def training_step(self, batch, batch_nb, optimizer_i):
+        imgs, _ = batch
+
+        # train generator
+        if optimizer_i == 0:
+            # sample noise
+            z = torch.randn(imgs.shape[0], self.hparams.latent_dim)
+
+            # match gpu device (or keep as cpu)
+            if self.on_gpu:
+                z = z.cuda(imgs.device.index)
+
+            # generate images
+            self.generated_imgs = self.forward(z)
+
+            # log sampled images
+            sample_imgs = self.generated_imgs[:6]
+            grid = torchvision.utils.make_grid(sample_imgs)
+            self.experiment.add_image('generated_images', grid, 0)
+
+            # ground truth result (ie: all fake)
+            valid = torch.ones(imgs.size(0), 1)
+
+            # adversarial loss is binary cross-entropy
+            g_loss = self.adversarial_loss(self.discriminator(self.generated_imgs), valid)
+
+            return g_loss
+
+        # train discriminator
+        if optimizer_i == 1:
+            # Measure discriminator's ability to classify real from generated samples
+
+            # how well can it label as real?
+            valid = torch.ones(imgs.size(0), 1)
+            real_loss = self.adversarial_loss(self.discriminator(imgs), valid)
+
+            # how well can it label as fake?
+            fake = torch.zeros(imgs.size(0), 1)
+            fake_loss = self.adversarial_loss(self.discriminator(self.generated_imgs.detach()), fake)
+
+            # discriminator loss is the average of these
+            d_loss = (real_loss + fake_loss) / 2
+
+            return d_loss
+
+    def configure_optimizers(self):
+        lr = self.hparams.lr
+        b1 = self.hparams.b1
+        b2 = self.hparams.b2
+
+        opt_g = torch.optim.Adam(self.generator.parameters(), lr=lr, betas=(b1, b2))
+        opt_d = torch.optim.Adam(self.discriminator.parameters(), lr=lr, betas=(b1, b2))
+        return [opt_g, opt_d], []
+
+    @pl.data_loader
+    def tng_dataloader(self):
+        transform = transforms.Compose([transforms.ToTensor(),
+                                        transforms.Normalize([0.5], [0.5])])
+        dataset = MNIST(os.getcwd(), train=True, download=True, transform=transform)
+        return DataLoader(dataset, batch_size=self.hparams.batch_size)
+
+
+def main(hparams):
+    # save tensorboard logs
+    exp = Experiment(save_dir=os.getcwd())
+
+    # init model
+    model = GAN(hparams)
+
+    # fit trainer on CPU
+    trainer = pl.Trainer(experiment=exp, max_nb_epochs=200)
+    trainer.fit(model)
+
+
+if __name__ == '__main__':
+    parser = ArgumentParser()
+    parser.add_argument("--batch_size", type=int, default=64, help="size of the batches")
+    parser.add_argument("--lr", type=float, default=0.0002, help="adam: learning rate")
+    parser.add_argument("--b1", type=float, default=0.5, help="adam: decay of first order momentum of gradient")
+    parser.add_argument("--b2", type=float, default=0.999, help="adam: decay of first order momentum of gradient")
+    parser.add_argument("--latent_dim", type=int, default=100, help="dimensionality of the latent space")
+
+    hparams = parser.parse_args()
+
+    main(hparams)

mkdocs.yml

@@ -8,3 +8,8 @@ site_description: 'Documentation for PyTorch LightningModule, the researcher ver
 
 dev_addr: '0.0.0.0:8000'
 #google_analytics: ['UA-aasd', 'sitename']
+
+markdown_extensions:
+  - codehilite:
+      guess_lang: false
+      linenums: true

pytorch_lightning/models/trainer.py

@@ -12,6 +12,7 @@ import torch
 from torch.utils.data.distributed import DistributedSampler
 import torch.multiprocessing as mp
 import torch.distributed as dist
+from torch.optim.optimizer import Optimizer
 
 from pytorch_lightning.root_module.root_module import LightningModule
 from pytorch_lightning.root_module.memory import get_gpu_memory_map
@@ -377,6 +378,31 @@ class Trainer(TrainerIO):
 
             self.tqdm_metrics[k] = v
 
+    def __validation_forward(self, model, data_batch, batch_i, dataloader_i):
+        # make dataloader_i arg in validation_step optional
+        args = [data_batch, batch_i]
+        if len(self.val_dataloader) > 1:
+            args.append(dataloader_i)
+
+        if self.use_ddp:
+            output = model(*args)
+        elif self.use_dp:
+            output = model(*args)
+        elif self.single_gpu:
+            # put inputs on gpu manually
+            gpu_id = self.data_parallel_device_ids[0]
+            data_batch = self.transfer_batch_to_gpu(data_batch, gpu_id)
+            args[0] = data_batch
+
+            # do non dp, ddp step
+            output = model.validation_step(*args)
+
+        else:
+            # CPU
+            output = model.validation_step(*args)
+
+        return output
+
     def validate(self, model, dataloader, max_batches, dataloader_i):
         """
         Run validation code
@@ -409,23 +435,9 @@ class Trainer(TrainerIO):
             # -----------------
             # RUN VALIDATION STEP
             # -----------------
-            if self.use_ddp:
-                output = model(data_batch, batch_i, dataloader_i)
-            elif self.use_dp:
-                output = model(data_batch, batch_i, dataloader_i)
-            elif self.single_gpu:
-                # put inputs on gpu manually
-                gpu_id = self.data_parallel_device_ids[0]
-                for i, x in enumerate(data_batch):
-                    if isinstance(x, torch.Tensor):
-                        data_batch[i] = x.cuda(gpu_id)
-
-                # do non dp, ddp step
-                output = model.validation_step(data_batch, batch_i, dataloader_i)
-
-            else:
-                output = model.validation_step(data_batch, batch_i, dataloader_i)
+            output = self.__validation_forward(model, data_batch, batch_i, dataloader_i)
 
+            # track outputs for collation
             outputs.append(output)
 
             # batch done
@@ -521,12 +533,15 @@ If you want each process to load the full dataset, ignore this warning.
                 task = int(os.environ['SLURM_LOCALID'])
                 self.ddp_train(task, model)
             else:
-                msg = """
-You requested %(nb_gpus)s GPUs but launched %(nb_tasks)s slurm tasks.
-We will launch %(nb_gpus)s processes for you.
-We recommend you let slurm manage the processes by setting: --ntasks-per-node=%(nb_gpus)s
-If you're not using SLURM, ignore this message!
-""" % {'nb_gpus': self.nb_requested_gpus, 'nb_tasks': self.nb_slurm_tasks}
+                nb_gpus = self.nb_requested_gpus
+                nb_tasks = self.nb_slurm_tasks
+                msg = f"""
+                You requested {nb_gpus}s GPUs but launched {nb_tasks}s slurm tasks.
+                We will launch {nb_gpus}s processes for you.
+                We recommend you let slurm manage the processes by setting:
+                --ntasks-per-node={nb_gpus}s
+                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, ))
 
@@ -547,9 +562,7 @@ If you're not using SLURM, ignore this message!
 
             # CHOOSE OPTIMIZER
             # allow for lr schedulers as well
-            self.optimizers = model.configure_optimizers()
-            if len(self.optimizers) == 2:
-                self.optimizers, self.lr_schedulers = self.optimizers
+            self.optimizers, self.lr_schedulers = self.init_optimizers(model.configure_optimizers())
 
             self.__run_pretrain_routine(model)
 
@@ -557,12 +570,25 @@ If you're not using SLURM, ignore this message!
         # used for testing or when we need to know that training succeeded
         return 1
 
+    def init_optimizers(self, optimizers):
+
+        # single optimizer
+        if isinstance(optimizers, Optimizer):
+            return [optimizers], []
+
+        # two lists
+        elif len(optimizers) == 2 and isinstance(optimizers[0], list):
+            optimizers, lr_schedulers = optimizers
+            return optimizers, lr_schedulers
+
+        # single list or tuple
+        elif isinstance(optimizers, list) or isinstance(optimizers, tuple):
+            return optimizers, []
+
     def __single_gpu_train(self, model):
         # CHOOSE OPTIMIZER
         # allow for lr schedulers as well
-        self.optimizers = model.configure_optimizers()
-        if len(self.optimizers) == 2:
-            self.optimizers, self.lr_schedulers = self.optimizers
+        self.optimizers, self.lr_schedulers = self.init_optimizers(model.configure_optimizers())
 
         model.cuda(self.data_parallel_device_ids[0])
 
@@ -579,20 +605,18 @@ If you're not using SLURM, ignore this message!
 
         # CHOOSE OPTIMIZER
         # allow for lr schedulers as well
-        self.optimizers = model.configure_optimizers()
-        if len(self.optimizers) == 2:
-            self.optimizers, self.lr_schedulers = self.optimizers
+        self.optimizers, self.lr_schedulers = self.init_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 = """
-Amp level %r with DataParallel is not supported.
-See this note from NVIDIA for more info: https://github.com/NVIDIA/apex/issues/227.
-We recommend you switch to ddp if you want to use amp
-""" % self.amp_level
+            m = f"""
+            Amp level {self.amp_level} with DataParallel is not supported.
+            See this note from NVIDIA for more info: https://github.com/NVIDIA/apex/issues/227.
+            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)
@@ -639,9 +663,7 @@ We recommend you switch to ddp if you want to use amp
 
         # CHOOSE OPTIMIZER
         # allow for lr schedulers as well
-        self.optimizers = model.configure_optimizers()
-        if len(self.optimizers) == 2:
-            self.optimizers, self.lr_schedulers = self.optimizers
+        self.optimizers, self.lr_schedulers = self.init_optimizers(model.configure_optimizers())
 
         # MODEL
         # copy model to each gpu
@@ -893,6 +915,95 @@ We recommend you switch to ddp if you want to use amp
         blacklist = {'batch_nb', 'v_nb', 'gpu'}
         return blacklist
 
+    def transfer_batch_to_gpu(self, batch, gpu_id):
+        # base case
+        if isinstance(batch, torch.Tensor):
+            return batch.cuda(gpu_id)
+
+        # when list
+        elif isinstance(batch, list):
+            for i, x in enumerate(batch):
+                batch[i] = self.transfer_batch_to_gpu(x, gpu_id)
+            return batch
+
+        # when dict
+        elif isinstance(batch, dict):
+            for k, v in batch.items():
+                batch[k] = self.transfer_batch_to_gpu(v, gpu_id)
+
+            return batch
+
+    def __tng_forward(self, data_batch, batch_nb, opt_idx):
+        """
+        Handle forward for each training case (distributed, single gpu, etc...)
+        :param data_batch:
+        :param batch_nb:
+        :return:
+        """
+        # ---------------
+        # FORWARD
+        # ---------------
+        # enable not needing to add opt_idx to training_step
+        args = [data_batch, batch_nb]
+        if len(self.optimizers) > 1:
+            args.append(opt_idx)
+
+        if self.use_ddp:
+            output = self.model(*args)
+        elif self.use_dp:
+            output = self.model(*args)
+        elif self.single_gpu:
+            gpu_id = self.data_parallel_device_ids[0]
+            data_batch = self.transfer_batch_to_gpu(data_batch, gpu_id)
+            args[0] = data_batch
+            output = self.model.training_step(*args)
+
+        else:
+            output = self.model.training_step(*args)
+
+        # ---------------
+        # TQDM metrics
+        # ---------------
+        try:
+            prog_output = output['prog']
+
+            # reduce prog metrics for tqdm when using dp
+            if self.use_dp:
+                nb_gpus = len(self.data_parallel_device_ids)
+                prog_output = reduce_distributed_output(prog_output, nb_gpus)
+
+            model_specific_tqdm_metrics_dic = prog_output
+        except Exception:
+            model_specific_tqdm_metrics_dic = {}
+
+        # ---------------
+        # EXTRACT LOSS
+        # ---------------
+        # if output dict doesn't have the keyword loss
+        # then assume the output=loss if scalar
+        try:
+            loss = output['loss']
+        except Exception:
+            if type(output) is torch.Tensor:
+                loss = output
+
+        # when using dp need to reduce the loss
+        if self.use_dp:
+            loss = reduce_distributed_output(loss, len(self.data_parallel_device_ids))
+
+        return loss, model_specific_tqdm_metrics_dic
+
+    def __clip_gradients(self):
+        if self.gradient_clip > 0:
+            model = self.__get_model()
+            torch.nn.utils.clip_grad_norm_(model.parameters(), self.gradient_clip)
+
+    def __print_nan_grads(self):
+        if self.print_nan_grads:
+            model = self.__get_model()
+            for param in model.parameters():
+                print(param.grad.float().sum())
+
     def __run_tng_batch(self, data_batch, batch_nb):
         if data_batch is None:
             return 0
@@ -908,102 +1019,57 @@ We recommend you switch to ddp if you want to use amp
         if self.progress_bar:
             self.prog_bar.update(1)
 
-        # forward pass
-        # return a scalar value and a dic with tqdm metrics
-        if self.use_ddp:
-            output = self.model(data_batch, batch_nb)
-        elif self.use_dp:
-            output = self.model(data_batch, batch_nb)
-        elif self.single_gpu:
-            gpu_id = self.data_parallel_device_ids[0]
-            for i, x in enumerate(data_batch):
-                if isinstance(x, torch.Tensor):
-                    data_batch[i] = x.cuda(gpu_id)
-            output = self.model.training_step(data_batch, batch_nb)
+        # call training_step once per optimizer
+        for opt_idx, optimizer in enumerate(self.optimizers):
 
-        else:
-            output = self.model.training_step(data_batch, batch_nb)
+            # forward pass
+            loss, model_specific_tqdm_metrics = self.__tng_forward(data_batch, batch_nb, opt_idx)
 
-        try:
-            prog_output = output['prog']
+            # track metrics
+            self.__add_tqdm_metrics(model_specific_tqdm_metrics)
 
-            # reduce prog metrics for tqdm when using dp
-            if self.use_dp:
-                nb_gpus = len(self.data_parallel_device_ids)
-                prog_output = reduce_distributed_output(prog_output, nb_gpus)
+            # accumulate loss
+            # (if accumulate_grad_batches = 1 no effect)
+            loss = loss / self.accumulate_grad_batches
 
-            model_specific_tqdm_metrics_dic = prog_output
-        except Exception:
-            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:
-            if type(output) is torch.Tensor:
-                loss = output
-
-        # when using dp need to reduce the loss
-        if self.use_dp:
-            loss = reduce_distributed_output(loss, len(self.data_parallel_device_ids))
-
-        self.__add_tqdm_metrics(model_specific_tqdm_metrics_dic)
-
-        # accumulate loss (if accumulate_grad_batches = 1 no effect)
-        loss = loss / self.accumulate_grad_batches
-
-        # backward pass
-        if self.use_amp:
-            # scale loss when using amp
-            for optimizer in self.optimizers:
+            # backward pass
+            if self.use_amp:
                 with amp.scale_loss(loss, optimizer) as scaled_loss:
                     scaled_loss.backward()
-        else:
-            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()
+            # insert after step hook
+            if self.__is_function_implemented('on_after_backward'):
+                model_ref = self.__get_model()
+                model_ref.on_after_backward()
 
-        if self.print_nan_grads:
-            model = self.__get_model()
-            for param in model.parameters():
-                print(param.grad.float().sum())
+            # nan grads
+            self.__print_nan_grads()
 
-        # track total loss for logging (avoid mem leaks)
-        self.batch_loss_value += loss.item()
+            # track total loss for logging (avoid mem 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:
+            # gradient update with accumulated gradients
+            if (self.batch_nb + 1) % self.accumulate_grad_batches == 0:
+                # clip gradients
+                self.__clip_gradients()
+
+                # calls .step(), .zero_grad()
+                # override function to modify this behavior
                 model = self.__get_model()
-                torch.nn.utils.clip_grad_norm_(model.parameters(), self.gradient_clip)
+                model.optimizer_step(self.current_epoch, batch_nb, optimizer, opt_idx)
 
-            # update gradients across all optimizers
-            for optimizer in self.optimizers:
-                optimizer.step()
+                # calculate running loss for display
+                self.running_loss.append(self.batch_loss_value)
+                self.batch_loss_value = 0
+                self.avg_loss = np.mean(self.running_loss[-100:])
 
-                # 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()
-
-            # calculate running loss for display
-            self.running_loss.append(self.batch_loss_value)
-            self.batch_loss_value = 0
-            self.avg_loss = np.mean(self.running_loss[-100:])
-
-            # update progbar
-            if self.progress_bar:
-                # add model specific metrics
-                tqdm_metrics = self.__tng_tqdm_dic
-                self.prog_bar.set_postfix(**tqdm_metrics)
+                # update progbar
+                if self.progress_bar:
+                    # add model specific metrics
+                    tqdm_metrics = self.__tng_tqdm_dic
+                    self.prog_bar.set_postfix(**tqdm_metrics)
 
         # activate batch end hook
         if self.__is_function_implemented('on_batch_end'):

pytorch_lightning/root_module/root_module.py

@@ -33,11 +33,21 @@ class LightningModule(GradInformation, ModelIO, ModelHooks):
         """
         raise NotImplementedError
 
-    def validation_step(self, data_batch, batch_nb):
+    def training_step(self, *args, **kwargs):
+        """
+        return loss, dict with metrics for tqdm
+        :param called with batch, batch_nb
+        additional: optimizer_i if multiple optimizers used
+        :return:
+        """
+        raise NotImplementedError
+
+    def validation_step(self, *args, **kwargs):
         """
         return whatever outputs will need to be aggregated in validation_end
         OPTIONAL
-        :param data_batch:
+        :param called with batch, batch_nb
+        additional: dataset_i if multiple val datasets used
         :return:
         """
         pass
@@ -51,14 +61,6 @@ class LightningModule(GradInformation, ModelIO, ModelHooks):
         """
         pass
 
-    def training_step(self, data_batch, batch_nb):
-        """
-        return loss, dict with metrics for tqdm
-        :param data_batch:
-        :return:
-        """
-        raise NotImplementedError
-
     def configure_optimizers(self):
         """
         Return a list of optimizers and a list of schedulers (could be empty)
@@ -66,6 +68,20 @@ class LightningModule(GradInformation, ModelIO, ModelHooks):
         """
         raise NotImplementedError
 
+    def optimizer_step(self, epoch_nb, batch_nb, optimizer, optimizer_i):
+        """
+        Do something instead of the standard optimizer behavior
+        :param epoch_nb:
+        :param batch_nb:
+        :param optimizer:
+        :param optimizer_i:
+        :return:
+        """
+        optimizer.step()
+
+        # clear gradients
+        optimizer.zero_grad()
+
     @data_loader
     def tng_dataloader(self):
         """

pytorch_lightning/testing/lm_test_module.py

@@ -195,7 +195,7 @@ class LightningTestModel(LightningModule):
         optimizer = optim.Adam(self.parameters(), lr=self.hparams.learning_rate)
 
         # test returning only 1 list instead of 2
-        return [optimizer]
+        return optimizer
 
     def __dataloader(self, train):
         # init data generators
@@ -231,7 +231,7 @@ class LightningTestModel(LightningModule):
 
     @data_loader
     def val_dataloader(self):
-        return self.__dataloader(train=False)
+        return [self.__dataloader(train=False), self.__dataloader(train=False)]
 
     @data_loader
     def test_dataloader(self):

pytorch_lightning/testing/no_val_end_module.py

@@ -109,7 +109,7 @@ class NoValEndTestModel(LightningModule):
         if self.trainer.batch_nb % 2 == 0:
             return loss_val
 
-    def validation_step(self, data_batch, batch_i, dataloader_i):
+    def validation_step(self, data_batch, batch_nb):
         """
         Lightning calls this inside the validation loop
         :param data_batch:
@@ -135,16 +135,16 @@ class NoValEndTestModel(LightningModule):
             val_acc = val_acc.unsqueeze(0)
 
         # alternate possible outputs to test
-        if batch_i % 1 == 0:
+        if batch_nb % 1 == 0:
             output = OrderedDict({
                 'val_loss': loss_val,
                 'val_acc': val_acc,
             })
             return output
-        if batch_i % 2 == 0:
+        if batch_nb % 2 == 0:
             return val_acc
 
-        if batch_i % 3 == 0:
+        if batch_nb % 3 == 0:
             output = OrderedDict({
                 'val_loss': loss_val,
                 'val_acc': val_acc,

setup.cfg

@@ -45,6 +45,7 @@ omit =
     tests/test_models.py
     pytorch_lightning/testing_models/lm_test_module.py
     pytorch_lightning/utilities/arg_parse.py
+    examples/templates
 
 [flake8]
 ignore = E731,W504,F401,F841

setup.py

@@ -14,7 +14,7 @@ from setuptools import setup, find_packages
 # engineer specific practices
 setup(
     name='pytorch-lightning',
-    version='0.4.4',
+    version='0.4.6',
     description='The Keras for ML researchers using PyTorch',
     author='William Falcon',
     author_email='waf2107@columbia.edu',

tests/test_models.py

@@ -27,6 +27,72 @@ np.random.seed(SEED)
 # TESTS
 # ------------------------------------------------------------------------
 
+def test_optimizer_return_options():
+
+    trainer = Trainer()
+    model, hparams = get_model()
+
+    # single optimizer
+    opt_a = torch.optim.Adam(model.parameters(), lr=0.002)
+    opt_b = torch.optim.SGD(model.parameters(), lr=0.002)
+    optim, lr_sched = trainer.init_optimizers(opt_a)
+    assert len(optim) == 1 and len(lr_sched) == 0
+
+    # opt tuple
+    opts = (opt_a, opt_b)
+    optim, lr_sched = trainer.init_optimizers(opts)
+    assert len(optim) == 2 and optim[0] == opts[0] and optim[1] == opts[1]
+    assert len(lr_sched) == 0
+
+    # opt list
+    opts = [opt_a, opt_b]
+    optim, lr_sched = trainer.init_optimizers(opts)
+    assert len(optim) == 2 and optim[0] == opts[0] and optim[1] == opts[1]
+    assert len(lr_sched) == 0
+
+    # opt tuple of lists
+    opts = ([opt_a], ['lr_scheduler'])
+    optim, lr_sched = trainer.init_optimizers(opts)
+    assert len(optim) == 1 and len(lr_sched) == 1
+    assert optim[0] == opts[0][0] and lr_sched[0] == 'lr_scheduler'
+
+
+def test_single_gpu_batch_parse():
+    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
+
+    trainer = Trainer()
+
+    # batch is just a tensor
+    batch = torch.rand(2, 3)
+    batch = trainer.transfer_batch_to_gpu(batch, 0)
+    assert batch.device.index == 0 and batch.type() == 'torch.cuda.FloatTensor'
+
+    # tensor list
+    batch = [torch.rand(2, 3), torch.rand(2, 3)]
+    batch = trainer.transfer_batch_to_gpu(batch, 0)
+    assert batch[0].device.index == 0 and batch[0].type() == 'torch.cuda.FloatTensor'
+    assert batch[1].device.index == 0 and batch[1].type() == 'torch.cuda.FloatTensor'
+
+    # tensor list of lists
+    batch = [[torch.rand(2, 3), torch.rand(2, 3)]]
+    batch = trainer.transfer_batch_to_gpu(batch, 0)
+    assert batch[0][0].device.index == 0 and batch[0][0].type() == 'torch.cuda.FloatTensor'
+    assert batch[0][1].device.index == 0 and batch[0][1].type() == 'torch.cuda.FloatTensor'
+
+    # tensor dict
+    batch = [{'a': torch.rand(2, 3), 'b': torch.rand(2, 3)}]
+    batch = trainer.transfer_batch_to_gpu(batch, 0)
+    assert batch[0]['a'].device.index == 0 and batch[0]['a'].type() == 'torch.cuda.FloatTensor'
+    assert batch[0]['b'].device.index == 0 and batch[0]['b'].type() == 'torch.cuda.FloatTensor'
+
+
 def test_early_stopping_cpu_model():
     """
     Test each of the trainer options
@@ -767,7 +833,7 @@ def test_multiple_val_dataloader():
     :return:
     """
     hparams = get_hparams()
-    model = LightningTemplateModel(hparams)
+    model = LightningTestModel(hparams)
 
     save_dir = init_save_dir()
 

tox.ini

@@ -38,7 +38,7 @@ commands =
     flake8 .
 
 [flake8]
-exclude = .tox,*.egg,build,temp
+exclude = .tox,*.egg,build,temp,examples/templates
 select = E,W,F
 doctests = True
 verbose = 2