release v0.4.5

* made dataset_i only available with multiple datasets

* updated interface signature

* updated tests

.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

README.md

@@ -91,7 +91,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
@@ -137,7 +137,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

docs/LightningModule/RequiredTrainerInterface.md

@@ -215,7 +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  |
+|  dataloader_i | Integer displaying which dataloader this is (only if multiple val datasets used)  |
 
 **Return**   
 
@@ -226,6 +226,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 +247,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 +382,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,33 @@ 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()    
+```
+
 ---
 #### on_before_zero_grad
 Called in the training loop after taking an optimizer step and before zeroing grads.

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):

pytorch_lightning/models/trainer.py

@@ -377,6 +377,32 @@ 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]
+            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(*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
@@ -893,6 +905,78 @@ We recommend you switch to ddp if you want to use amp
         blacklist = {'batch_nb', 'v_nb', 'gpu'}
         return blacklist
 
+    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]
+            for i, x in enumerate(data_batch):
+                if isinstance(x, torch.Tensor):
+                    data_batch[i] = x.cuda(gpu_id)
+            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 +992,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

@@ -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.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.5',
     description='The Keras for ML researchers using PyTorch',
     author='William Falcon',
     author_email='waf2107@columbia.edu',

tests/test_models.py

@@ -767,7 +767,7 @@ def test_multiple_val_dataloader():
     :return:
     """
     hparams = get_hparams()
-    model = LightningTemplateModel(hparams)
+    model = LightningTestModel(hparams)
 
     save_dir = init_save_dir()