Support for multiple val_dataloaders (#97)

* Added support for multiple validation dataloaders

* Fix typo in README.md

* Update trainer.py

* Add support for multiple dataloaders

* Rename dataloader_index to dataloader_i

* Added warning to check val_dataloaders

Added a warning to ensure that all val_dataloaders were DistributedSamplers if ddp is enabled

* Updated DistributedSampler warning

* Fixed typo

* Added multiple val_dataloaders

* Multiple val_dataloader test

* Update lightning_module_template.py

Added dataloader_i to validation_step parameters

* Update trainer.py

* Reverted template changes

* Create multi_val_module.py

* Update no_val_end_module.py

* New MultiValModel

* Rename MultiValModel to MultiValTestModel

* Revert to LightningTestModel

* Update test_models.py

* Update trainer.py

* Update test_models.py

* multiple val_dataloaders in test template

* Fixed flake8 warnings

* Update trainer.py

* Fix flake errors

* Fixed Flake8 errors

* Update lm_test_module.py

keep this test model with a single dataset for val

* Update trainer.py

* Update trainer.py

* Update trainer.py

* Update trainer.py

* Update trainer.py

* Update test_models.py

* Update trainer.py

* Update trainer.py

* Update trainer.py

* Update trainer.py

* Update trainer.py

* Update trainer.py

* Update RequiredTrainerInterface.md

* Update RequiredTrainerInterface.md

* Update test_models.py

* Update trainer.py

dont need the else clause, val_dataloader is either a list or none because of get_dataloaders()

* Update trainer.py

fixed flake errors

* Update trainer.py

Author: sholalkere

docs/LightningModule/RequiredTrainerInterface.md

@@ -46,24 +46,25 @@ class CoolModel(pl.LightningModule):
     def forward(self, 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)
-
     def training_step(self, batch, batch_nb):
+        # REQUIRED
         x, y = batch
         y_hat = self.forward(x)
-        return {'loss': self.my_loss(y_hat, y)}
+        return {'loss': F.cross_entropy(y_hat, y)(y_hat, y)}
 
     def validation_step(self, batch, batch_nb):
+        # 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)(y_hat, y)}
 
     def validation_end(self, outputs):
+        # OPTIONAL
         avg_loss = torch.stack([x['val_loss'] for x in outputs]).mean()
         return {'avg_val_loss': avg_loss}
 
     def configure_optimizers(self):
+        # REQUIRED
         return [torch.optim.Adam(self.parameters(), lr=0.02)]
 
     @pl.data_loader
@@ -72,10 +73,13 @@ class CoolModel(pl.LightningModule):
 
     @pl.data_loader
     def val_dataloader(self):
+        # OPTIONAL
+        # can also return a list of val dataloaders
         return DataLoader(MNIST(os.getcwd(), train=True, download=True, transform=transforms.ToTensor()), batch_size=32)
 
     @pl.data_loader
     def test_dataloader(self):
+        # OPTIONAL
         return DataLoader(MNIST(os.getcwd(), train=True, download=True, transform=transforms.ToTensor()), batch_size=32)
 ```
 ---   
@@ -88,7 +92,7 @@ The LightningModule interface is on the right. Each method corresponds to a part
   </a>
 </p>   
 
----
+## Required Methods    
 
 ### training_step
 
@@ -134,15 +138,75 @@ def training_step(self, data_batch, batch_nb):
     return output
 ```
 
----
+--- 
+### tng_dataloader 
+
+``` {.python}
+@pl.data_loader
+def tng_dataloader(self)
+```
+Called by lightning during training loop. Make sure to use the @pl.data_loader decorator, this ensures not calling this function until the data are needed.
+
+##### Return
+PyTorch DataLoader
+
+**Example**
+
+``` {.python}
+@pl.data_loader
+def tng_dataloader(self):
+    transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,), (1.0,))])
+    dataset = MNIST(root='/path/to/mnist/', train=True, transform=transform, download=True)
+    loader = torch.utils.data.DataLoader(
+        dataset=dataset,
+        batch_size=self.hparams.batch_size,
+        shuffle=True
+    )
+    return loader
+```
+
+--- 
+### configure_optimizers 
+
+``` {.python}
+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.
+
+
+##### Return
+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(self.parameters(), lr=0.01)
+    return [opt]
+    
+# gan example, with scheduler for discriminator
+def configure_optimizers(self):
+    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]
+```
+
+## Optional Methods    
 
 ### validation_step
 
 ``` {.python}
-def validation_step(self, data_batch, batch_nb)
+def validation_step(self, data_batch, batch_nb, dataloader_i)
 ```
+**OPTIONAL**    
+If you don't need to validate you don't need to implement this method.    
+
+In this step you'd normally do the forward pass and calculate the loss for a batch. You can also do fancier things like multiple forward passes, calculate accuracy, or save example outputs (using self.experiment or whatever you want). Really, anything you want.       
 
-In this step you'd normally do the forward pass and calculate the loss for a batch. You can also do fancier things like multiple forward passes or something specific to your model.
 This is most likely the same as your training_step. But unlike training step, the outputs from here will go to validation_end for collation.
 
 **Params**   
@@ -151,6 +215,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  |
 
 **Return**   
 
@@ -188,9 +253,12 @@ def validation_step(self, data_batch, batch_nb):
 
 ``` {.python}
 def validation_end(self, outputs)
-```
+```   
+If you didn't define a validation_step, this won't be called.       
+
+Called at the end of the validation loop with the output of each validation_step.  Called once per validation dataset.   
 
-Called at the end of the validation loop with the output of each validation_step.
+The outputs here are strictly for the progress bar. If you don't need to display anything, don't return anything.    
 
 **Params**    
 
@@ -225,36 +293,6 @@ def validation_end(self, outputs):
     return tqdm_dic
 ```
 
---- 
-### configure_optimizers 
-
-``` {.python}
-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.
-
-
-##### Return
-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(self.parameters(), lr=0.01)
-    return [opt]
-    
-# gan example, with scheduler for discriminator
-def configure_optimizers(self):
-    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]
-```
-
 --- 
 ### on_save_checkpoint 
 
@@ -297,44 +335,20 @@ def on_load_checkpoint(self, checkpoint):
     self.something_cool_i_want_to_save = checkpoint['something_cool_i_want_to_save']
 ```
 
---- 
-### tng_dataloader 
-
-``` {.python}
-@pl.data_loader
-def tng_dataloader(self)
-```
-Called by lightning during training loop. Make sure to use the @pl.data_loader decorator, this ensures not calling this function until the data are needed.
-
-##### Return
-PyTorch DataLoader
-
-**Example**
-
-``` {.python}
-@pl.data_loader
-def tng_dataloader(self):
-    transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,), (1.0,))])
-    dataset = MNIST(root='/path/to/mnist/', train=True, transform=transform, download=True)
-    loader = torch.utils.data.DataLoader(
-        dataset=dataset,
-        batch_size=self.hparams.batch_size,
-        shuffle=True
-    )
-    return loader
-```
-
 --- 
 ### val_dataloader 
 
 ``` {.python}
 @pl.data_loader
 def tng_dataloader(self)
 ```
-Called by lightning during validation loop. Make sure to use the @pl.data_loader decorator, this ensures not calling this function until the data are needed.
+**OPTIONAL**    
+If you don't need a validation dataset and a validation_step, you don't need to implement this method.    
+
+Called by lightning during validation loop. Make sure to use the @pl.data_loader decorator, this ensures not calling this function until the data are needed.   
 
 ##### Return
-PyTorch DataLoader
+PyTorch DataLoader or list of PyTorch Dataloaders.    
 
 **Example**
 
@@ -350,6 +364,11 @@ def val_dataloader(self):
     )
     
     return loader
+
+# can also return multiple dataloaders   
+@pl.data_loader
+def val_dataloader(self):
+    return [loader_a, loader_b, ..., loader_n]   
 ```
 
 --- 
@@ -359,6 +378,9 @@ def val_dataloader(self):
 @pl.data_loader
 def test_dataloader(self)
 ```
+**OPTIONAL**    
+If you don't need a test dataset and a test_step, you don't need to implement this method.    
+
 Called by lightning during test loop. Make sure to use the @pl.data_loader decorator, this ensures not calling this function until the data are needed.
 
 ##### Return

examples/new_project_templates/lightning_module_template.py

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