Expectopatronum implement #89 (#182)

* rename validate -> evaluate; implement test logic; allow multiple test_loaders

* add test_step and test_end to LightningModule

* add in_test_mode to pretraining to implement case 2 (test pretrained model)

* fix code style issues

* LightningTestModel: add optional second test set, implement test_step and test_end

* implemented test for multiple test_dataloaders; fixed typo

* add two test cases for #89

* add documentation for test_step, test_end; fix computation of loss in validation_step example

* Update trainer.py

* Update trainer.py

* Update trainer.py

* Update trainer.py

* Update trainer.py

* Update trainer.py

* Added proper dp ddp routing calls for test mode

* Update trainer.py

* Update test_models.py

* Update trainer.py

* Update trainer.py

* Update override_data_parallel.py

* Update test_models.py

* Update test_models.py

* Update trainer.py

* Update trainer.py

* Update trainer.py

* Update test_models.py

* Update test_models.py

* debug

* debug

* debug

* debug

* debug

* debug

* debug

* debug

* debug

* debug

* debug

* debug

* debug

* debug

* debug

* debug

* debug

* debug

* debug

* debug

* debug

* debug

* debug

* debug

* debug

* debug

* debug

* debug

* debug

* debug

* debug

* debug

* debug

* debug

* debug

* debug

* debug

* debug

* debug

* debug

* debug

* debug

* debug

* debug

* debug

* debug

* debug

* debug

* debug

* debug

* debug

* debug

* debug

* debug

* debug

* debug

* debug

* debug

* debug

* debug

* debug

* debug

* debug

* debug

* debug

* debug

* debug

* debug

* debug

* debug

* debug

* debug

* debug

* debug

* debug

* debug

* debug

* debug

* debug

* debug

* debug

* debug

* debug

* debug

* debug

* debug

* debug

* debug

* debug

* debug

* debug

* Update trainer.py

* Update override_data_parallel.py

* Update debug.py

* Update lm_test_module.py

* Update test_models.py

docs/LightningModule/RequiredTrainerInterface.md

@@ -17,6 +17,8 @@ Otherwise, to Define a Lightning Module, implement the following methods:
 
 - [validation_step](RequiredTrainerInterface.md#validation_step)    
 - [validation_end](RequiredTrainerInterface.md#validation_end) 
+- [test_step](RequiredTrainerInterface.md#test_step)    
+- [test_end](RequiredTrainerInterface.md#test_end) 
 - [val_dataloader](RequiredTrainerInterface.md#val_dataloader)    
 - [test_dataloader](RequiredTrainerInterface.md#test_dataloader)    
 - [on_save_checkpoint](RequiredTrainerInterface.md#on_save_checkpoint)    
@@ -63,6 +65,17 @@ class CoolModel(pl.LightningModule):
         avg_loss = torch.stack([x['val_loss'] for x in outputs]).mean()
         return {'avg_val_loss': avg_loss}
 
+    def test_step(self, batch, batch_nb):
+        # OPTIONAL
+        x, y = batch
+        y_hat = self.forward(x)
+        return {'test_loss': F.cross_entropy(y_hat, y)}
+
+    def test_end(self, outputs):
+        # OPTIONAL
+        avg_loss = torch.stack([x['test_loss'] for x in outputs]).mean()
+        return {'avg_test_loss': avg_loss}
+
     def configure_optimizers(self):
         # REQUIRED
         return [torch.optim.Adam(self.parameters(), lr=0.02)]
@@ -80,6 +93,7 @@ class CoolModel(pl.LightningModule):
     @pl.data_loader
     def test_dataloader(self):
         # OPTIONAL
+        # can also return a list of test dataloaders
         return DataLoader(MNIST(os.getcwd(), train=False, download=True, transform=transforms.ToTensor()), batch_size=32)
 ```
 ---   
@@ -225,9 +239,10 @@ the [optimizer_step](https://williamfalcon.github.io/pytorch-lightning/Trainer/h
 ### validation_step
 
 ``` {.python}
+# if you have one val dataloader:
 def validation_step(self, data_batch, batch_nb)   
 
-# if have multiple val dataloaders:  
+# if you have multiple val dataloaders:  
 def validation_step(self, data_batch, batch_nb, dataloader_idx)
 ```
 **OPTIONAL**    
@@ -235,7 +250,7 @@ If you don't need to validate you don't need to implement this method.
 
 In this step you'd normally generate examples or calculate anything of interest such as accuracy. 
 
-The dict you return here will be available in the validation_end method. 
+The dict you return here will be available in the `validation_end` method. 
 
 **Params**   
 
@@ -256,11 +271,11 @@ The dict you return here will be available in the validation_end method.
 ``` {.python}
 # CASE 1: A single validation dataset
 def validation_step(self, data_batch, batch_nb):
-    x, y, z = data_batch
+    x, y = data_batch
     
     # implement your own
     out = self.forward(x)
-    loss = self.loss(out, x)
+    loss = self.loss(out, y)
     
     # log 6 example images
     # or generated text... or whatever
@@ -338,6 +353,119 @@ def validation_end(self, outputs):
     return tqdm_dic
 ```
 
+### test_step
+
+``` {.python}
+# if you have one test dataloader:
+def test_step(self, data_batch, batch_nb)   
+
+# if you have multiple test dataloaders:  
+def test_step(self, data_batch, batch_nb, dataloader_idx)
+```
+**OPTIONAL**    
+If you don't need to test you don't need to implement this method.    
+
+In this step you'd normally generate examples or calculate anything of interest such as accuracy. 
+
+The dict you return here will be available in the `test_end` method. 
+
+This function is used when you execute `trainer.test()`.
+
+**Params**   
+
+| Param  | description  |
+|---|---|
+|  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 (only if multiple test datasets used)  |
+
+**Return**   
+
+| Return  | description  | optional |
+|---|---|---|   
+|  dict | Dict or OrderedDict with metrics to display in progress bar. All keys must be tensors. | Y |
+
+**Example**
+
+``` {.python}
+# CASE 1: A single test dataset
+def test_step(self, data_batch, batch_nb):
+    x, y = data_batch
+    
+    # implement your own
+    out = self.forward(x)
+    loss = self.loss(out, y)
+    
+    # calculate acc
+    labels_hat = torch.argmax(out, dim=1)
+    test_acc = torch.sum(y == labels_hat).item() / (len(y) * 1.0)
+    
+    # all optional...
+    # return whatever you need for the collation function test_end
+    output = OrderedDict({
+        'test_loss': loss_test,
+        'test_acc': torch.tensor(test_acc), # everything must be a tensor
+    })
+    
+    # return an optional dict
+    return output
+```   
+
+If you pass in multiple test datasets, test_step will have an additional argument.
+
+```python
+# CASE 2: multiple test datasets
+def test_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 ```test_dataloader```.    
+
+--- 
+### test_end
+
+``` {.python}
+def test_end(self, outputs)
+```   
+If you didn't define a test_step, this won't be called.       
+
+Called at the end of the test step with the output of each test_step.  Called once per test dataset.   
+
+The outputs here are strictly for the progress bar. If you don't need to display anything, don't return anything.    
+
+**Params**    
+
+| Param  | description  |
+|---|---|
+|  outputs | List of outputs you defined test_step |
+
+**Return**   
+
+| Return  | description  | optional |
+|---|---|---|   
+|  dict | Dict of OrderedDict with metrics to display in progress bar | Y |
+
+**Example**
+
+``` {.python}
+def test_end(self, outputs):
+    """
+    Called at the end of test to aggregate outputs
+    :param outputs: list of individual outputs of each test step
+    :return:
+    """
+    test_loss_mean = 0
+    test_acc_mean = 0
+    for output in outputs:
+        test_loss_mean += output['test_loss']
+        test_acc_mean += output['test_acc']
+
+    test_loss_mean /= len(outputs)
+    test_acc_mean /= len(outputs)
+    tqdm_dic = {'test_loss': test_loss_mean.item(), 'test_acc': test_acc_mean.item()}
+    return tqdm_dic
+```
+
 --- 
 ### on_save_checkpoint 
 

pytorch_lightning/models/trainer.py

@@ -20,6 +20,7 @@ from pytorch_lightning.pt_overrides.override_data_parallel import (
     LightningDistributedDataParallel, LightningDataParallel)
 from pytorch_lightning.callbacks import GradientAccumulationScheduler
 from pytorch_lightning.utilities.debugging import MisconfigurationException
+import pdb
 
 try:
     from apex import amp
@@ -157,6 +158,7 @@ class Trainer(TrainerIO):
         self.use_ddp = False
         self.use_dp = False
         self.single_gpu = False
+        self.testing = False
 
         # training bookeeping
         self.total_batch_nb = 0
@@ -370,8 +372,10 @@ class Trainer(TrainerIO):
             self.nb_val_batches = max(1, self.nb_val_batches)
 
         # determine number of test batches
-        self.nb_test_batches = len(self.test_dataloader) if self.test_dataloader is not None else 0
+        if self.test_dataloader is not None:
-        self.nb_test_batches = int(self.nb_test_batches * self.test_percent_check)
+            self.nb_test_batches = sum(len(dataloader) for dataloader in self.test_dataloader)
+            self.nb_test_batches = int(self.nb_test_batches * self.test_percent_check)
+            self.nb_test_batches = max(1, self.nb_test_batches)
 
         # determine when to check validation
         self.val_check_batch = int(self.nb_tng_batches * self.val_check_interval)
@@ -384,40 +388,45 @@ class Trainer(TrainerIO):
 
             self.tqdm_metrics[k] = v
 
-    def __validation_forward(self, model, data_batch, batch_i, dataloader_i):
+    def __evaluation_forward(self, model, data_batch, batch_i, dataloader_i, test=False):
         # make dataloader_i arg in validation_step optional
         args = [data_batch, batch_i]
-        if len(self.val_dataloader) > 1:
+
+        if test and len(self.test_dataloader) > 1:
             args.append(dataloader_i)
 
-        if self.use_ddp:
+        elif len(self.val_dataloader) > 1:
+            args.append(dataloader_i)
+
+        # handle DP, DDP forward
+        if self.use_ddp or self.use_dp:
             output = model(*args)
-        elif self.use_dp:
+            return output
-            output = model(*args)
+
-        elif self.single_gpu:
+        # CPU, single GPU
-            # put inputs on gpu manually
+        if self.single_gpu:
+            # for 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
+        if test:
-            output = model.validation_step(*args)
+            output = model.test_step(*args)
-
         else:
-            # CPU
             output = model.validation_step(*args)
 
         return output
 
-    def validate(self, model, dataloader, max_batches, dataloader_i):
+    def evaluate(self, model, dataloader, max_batches, dataloader_i, test=False):
         """
-        Run validation code
+        Run evaluation code
         :param model: PT model
         :param dataloader: PT dataloader
         :param max_batches: Scalar
+        :param dataloader_i:
+        :param test: boolean
         :return:
         """
-
         # enable eval mode
         model.zero_grad()
         model.eval()
@@ -439,9 +448,10 @@ class Trainer(TrainerIO):
                 break
 
             # -----------------
-            # RUN VALIDATION STEP
+            # RUN EVALUATION STEP
             # -----------------
-            output = self.__validation_forward(model, data_batch, batch_i, dataloader_i)
+            output = self.__evaluation_forward(model, data_batch, batch_i, dataloader_i,
+                                               test)
 
             # track outputs for collation
             outputs.append(output)
@@ -450,13 +460,14 @@ class Trainer(TrainerIO):
             if self.show_progress_bar:
                 self.progress_bar.update(1)
 
+        eval_results = {}
+
         # give model a chance to do something with the outputs (and method defined)
-        val_results = {}
+        model = self.__get_model()
-        if self.__is_overriden('validation_end'):
+        if test and self.__is_overriden('test_end'):
-            if self.data_parallel:
+            eval_results = model.test_end(outputs)
-                val_results = model.module.validation_end(outputs)
+        elif self.__is_overriden('validation_end'):
-            else:
+            eval_results = model.validation_end(outputs)
-                val_results = model.validation_end(outputs)
 
         # enable train mode again
         model.train()
@@ -464,7 +475,7 @@ class Trainer(TrainerIO):
         # enable gradients to save memory
         torch.set_grad_enabled(True)
 
-        return val_results
+        return eval_results
 
     def get_dataloaders(self, model):
         """
@@ -478,6 +489,10 @@ class Trainer(TrainerIO):
         self.val_dataloader = model.val_dataloader
 
         # handle returning an actual dataloader instead of a list of loaders
+        have_test_loaders = self.test_dataloader is not None
+        if have_test_loaders and not isinstance(self.test_dataloader, list):
+            self.test_dataloader = [self.test_dataloader]
+
         have_val_loaders = self.val_dataloader is not None
         if have_val_loaders and not isinstance(self.val_dataloader, list):
             self.val_dataloader = [self.val_dataloader]
@@ -524,11 +539,33 @@ class Trainer(TrainerIO):
                     warnings.warn(msg)
                     break
 
+        if self.use_ddp and self.test_dataloader is not None:
+            for dataloader in self.test_dataloader:
+                if not isinstance(dataloader, DistributedSampler):
+                    msg = """
+                    Your test_dataloader(s) are not all DistributedSamplers.
+                    You're using multiple gpus and multiple nodes without using a DistributedSampler
+                    to assign a subset of your data to each process. To silence this warning, pass a
+                    DistributedSampler to your DataLoader.
+
+                    ie: this:
+                    dataset = myDataset()
+                    dataloader = Dataloader(dataset)
+
+                    becomes:
+                    dataset = myDataset()
+                    dist_sampler = torch.utils.data.distributed.DistributedSampler(dataset)
+                    dataloader = Dataloader(dataset, sampler=dist_sampler)
+
+                    If you want each process to load the full dataset, ignore this warning.
+                    """
+                    warnings.warn(msg)
+                    break
+
     # -----------------------------
     # MODEL TRAINING
     # -----------------------------
     def fit(self, model):
-
         # when using multi-node or DDP within a node start each module in a separate process
         if self.use_ddp:
             # must copy only the meta of the exp so it survives pickle/unpickle
@@ -740,6 +777,7 @@ class Trainer(TrainerIO):
         if self.data_parallel:
             ref_model = model.module
 
+        # give model convenience properties
         ref_model.trainer = self
 
         # set local properties on the model
@@ -747,6 +785,7 @@ class Trainer(TrainerIO):
         ref_model.use_dp = self.use_dp
         ref_model.use_ddp = self.use_ddp
         ref_model.use_amp = self.use_amp
+        ref_model.testing = self.testing
 
         # transfer data loaders from model
         self.get_dataloaders(ref_model)
@@ -758,15 +797,13 @@ class Trainer(TrainerIO):
         if self.proc_rank == 0 and self.print_weights_summary:
             ref_model.summarize()
 
-        # give model convenience properties
+        # link up experiment object
-        ref_model.trainer = self
-
         if self.experiment is not None:
             ref_model.experiment = self.experiment
 
-        # save exp to get started
+            # save exp to get started
-        if self.proc_rank == 0 and self.experiment is not None:
+            if self.proc_rank == 0:
-            self.experiment.save()
+                self.experiment.save()
 
         # track model now.
         # if cluster resets state, the model will update with the saved weights
@@ -785,7 +822,13 @@ class Trainer(TrainerIO):
         if self.show_progress_bar:
             self.progress_bar = tqdm.tqdm(0, position=self.process_position)
 
-        # run tiny validation (if validation defined) to make sure program won't crash during val
+        # when testing requested only run test and return
+        if self.testing:
+            self.__run_evaluation(test=True)
+            return
+
+        # run tiny validation (if validation defined)
+        # to make sure program won't crash during val
         ref_model.on_sanity_check_start()
         if self.val_dataloader is not None and self.nb_sanity_val_steps > 0:
             for ds_i, dataloader in enumerate(self.val_dataloader):
@@ -794,7 +837,7 @@ class Trainer(TrainerIO):
                 if self.show_progress_bar:
                     self.progress_bar.reset(self.nb_sanity_val_steps)
 
-                self.validate(model, dataloader, self.nb_sanity_val_steps, ds_i)
+                self.evaluate(model, dataloader, self.nb_sanity_val_steps, ds_i, self.testing)
 
         # ---------------------------
         # CORE TRAINING LOOP
@@ -871,8 +914,10 @@ class Trainer(TrainerIO):
             # RUN VAL STEP
             # ---------------
             is_val_check_batch = (batch_nb + 1) % self.val_check_batch == 0
+            can_check_epoch = (self.current_epoch + 1) % self.check_val_every_n_epoch == 0
             if self.fast_dev_run or is_val_check_batch or early_stop_epoch:
-                self.__run_validation()
+                if can_check_epoch:
+                    self.__run_evaluation(test=self.testing)
 
             # when batch should be saved
             if (batch_nb + 1) % self.log_save_interval == 0 or early_stop_epoch:
@@ -917,6 +962,13 @@ class Trainer(TrainerIO):
             model = self.__get_model()
             model.on_epoch_end()
 
+    def test(self, model=None):
+        if model is not None:
+            self.testing = True
+            self.fit(model)
+        else:
+            self.__run_evaluation(test=True)
+
     def __metrics_to_scalars(self, metrics, blacklist=set()):
         new_metrics = {}
         for k, v in metrics.items():
@@ -1109,32 +1161,49 @@ class Trainer(TrainerIO):
 
         return 0
 
-    def __run_validation(self):
+    def __run_evaluation(self, test=False):
-        # decide if can check epochs
+        # when testing make sure user defined a test step
-        can_check_epoch = (self.current_epoch + 1) % self.check_val_every_n_epoch == 0
+        can_run_test_step = False
-        if self.fast_dev_run:
+        if test:
-            print('skipping to check performance bc of --fast_dev_run')
+            can_run_test_step = self.__is_overriden('test_step') and self.__is_overriden('test_end')
-        elif not can_check_epoch:
+            if not can_run_test_step:
-            return
+                m = '''You called .test() without defining a test step or test_end.
+                Please define and try again'''
+                raise MisconfigurationException(m)
 
         # validate only if model has validation_step defined
-        if self.__is_overriden('validation_step'):
+        # test only if test_step or validation_step are defined
+        run_val_step = self.__is_overriden('validation_step')
+
+        if run_val_step or can_run_test_step:
 
             # hook
-            if self.__is_function_implemented('on_pre_performance_check'):
+            model = self.__get_model()
-                model = self.__get_model()
+            model.on_pre_performance_check()
-                model.on_pre_performance_check()
 
-            # use val_percent_check set on end of epoch
+            # select dataloaders
-            # use a small portion otherwise
+            dataloaders = self.val_dataloader
-            max_batches = self.nb_val_batches if not self.fast_dev_run else 1
+            max_batches = self.nb_val_batches
-            for ds_i, dataloader in enumerate(self.val_dataloader):
-                val_out_metrics = self.validate(self.model, dataloader, max_batches, ds_i)
-                self.__add_tqdm_metrics(val_out_metrics)
 
-            # hook
+            # calculate max batches to use
-            if self.__is_function_implemented('on_post_performance_check'):
+            if test:
-                model = self.__get_model()
+                dataloaders = self.test_dataloader
+                max_batches = self.nb_test_batches
+
+            # cap max batches to 1 when using fast_dev_run
+            if self.fast_dev_run:
+                max_batches = 1
+
+            for ds_i, dataloader in enumerate(dataloaders):
+                eval_out_metrics = self.evaluate(self.model,
+                                                 dataloader,
+                                                 max_batches,
+                                                 ds_i,
+                                                 test)
+
+                self.__add_tqdm_metrics(eval_out_metrics)
+
+                # hook
                 model.on_post_performance_check()
 
             if self.show_progress_bar:
@@ -1143,7 +1212,7 @@ class Trainer(TrainerIO):
                 self.progress_bar.set_postfix(**tqdm_metrics)
 
         # model checkpointing
-        if self.proc_rank == 0 and self.checkpoint_callback is not None:
+        if self.proc_rank == 0 and self.checkpoint_callback is not None and not test:
             print('save callback...')
             self.checkpoint_callback.on_epoch_end(epoch=self.current_epoch,
                                                   logs=self.__tng_tqdm_dic)

pytorch_lightning/pt_overrides/override_data_parallel.py

@@ -56,6 +56,8 @@ class LightningDataParallel(DataParallel):
             # lightning
             if self.module.training:
                 return self.module.training_step(*inputs[0], **kwargs[0])
+            elif self.module.testing:
+                return self.module.test_step(*inputs[0], **kwargs[0])
             else:
                 return self.module.validation_step(*inputs[0], **kwargs[0])
 
@@ -89,6 +91,8 @@ class LightningDistributedDataParallel(DistributedDataParallel):
                 # lightning
                 if self.module.training:
                     output = self.module.training_step(*inputs[0], **kwargs[0])
+                elif self.module.testing:
+                    output = self.module.test_step(*inputs[0], **kwargs[0])
                 else:
                     output = self.module.validation_step(*inputs[0], **kwargs[0])
             else:
@@ -153,6 +157,10 @@ def parallel_apply(modules, inputs, kwargs_tup=None, devices=None):  # pragma: n
                 # CHANGE
                 if module.training:
                     output = module.training_step(*input, **kwargs)
+
+                elif module.testing:
+                    output = module.test_step(*input, **kwargs)
+
                 else:
                     output = module.validation_step(*input, **kwargs)
                 # ---------------

pytorch_lightning/root_module/hooks.py

@@ -5,7 +5,7 @@ class ModelHooks(torch.nn.Module):
 
     def on_sanity_check_start(self):
         """
-        Called before starting validate
+        Called before starting evaluate
         :return:
         """
         pass

pytorch_lightning/root_module/root_module.py

@@ -55,6 +55,16 @@ class LightningModule(GradInformation, ModelIO, ModelHooks):
         """
         pass
 
+    def test_step(self, *args, **kwargs):
+        """
+        return whatever outputs will need to be aggregated in test_end
+        OPTIONAL
+        :param called with batch, batch_nb
+        additional: dataset_i if multiple val datasets used
+        :return:
+        """
+        pass
+
     def validation_end(self, outputs):
         """
         Outputs has the appended output after each validation step
@@ -64,6 +74,15 @@ class LightningModule(GradInformation, ModelIO, ModelHooks):
         """
         pass
 
+    def test_end(self, outputs):
+        """
+        Outputs has the appended output after each test step
+        OPTIONAL
+        :param outputs:
+        :return: dic_with_metrics for tqdm
+        """
+        pass
+
     def configure_optimizers(self):
         """
         Return a list of optimizers and a list of schedulers (could be empty)

pytorch_lightning/testing/lm_test_module.py

@@ -20,7 +20,7 @@ class LightningTestModel(LightningModule):
     Sample model to show how to define a template
     """
 
-    def __init__(self, hparams, force_remove_distributed_sampler=False):
+    def __init__(self, hparams, force_remove_distributed_sampler=False, use_two_test_sets=False):
         """
         Pass in parsed HyperOptArgumentParser to the model
         :param hparams:
@@ -28,6 +28,7 @@ class LightningTestModel(LightningModule):
         # init superclass
         super(LightningTestModel, self).__init__()
         self.hparams = hparams
+        self.use_two_test_sets = use_two_test_sets  # for some tests regarding testing
 
         self.batch_size = hparams.batch_size
 
@@ -167,12 +168,22 @@ class LightningTestModel(LightningModule):
         # if returned a scalar from validation_step, outputs is a list of tensor scalars
         # we return just the average in this case (if we want)
         # return torch.stack(outputs).mean()
-
         val_loss_mean = 0
         val_acc_mean = 0
         for output in outputs:
-            val_loss_mean += output['val_loss']
+            val_loss = output['val_loss']
-            val_acc_mean += output['val_acc']
+
+            # reduce manually when using dp
+            if self.trainer.use_dp:
+                val_loss = torch.mean(val_loss)
+            val_loss_mean += val_loss
+
+            # reduce manually when using dp
+            val_acc = output['val_acc']
+            if self.trainer.use_dp:
+                val_acc = torch.mean(val_acc)
+
+            val_acc_mean += val_acc
 
         val_loss_mean /= len(outputs)
         val_acc_mean /= len(outputs)
@@ -180,6 +191,87 @@ class LightningTestModel(LightningModule):
         tqdm_dic = {'val_loss': val_loss_mean.item(), 'val_acc': val_acc_mean.item()}
         return tqdm_dic
 
+    def test_step(self, data_batch, batch_i, dataloader_i):
+        """
+        Lightning calls this 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_test = self.loss(y, y_hat)
+
+        # acc
+        labels_hat = torch.argmax(y_hat, dim=1)
+        test_acc = torch.sum(y == labels_hat).item() / (len(y) * 1.0)
+        test_acc = torch.tensor(test_acc)
+
+        if self.on_gpu:
+            test_acc = test_acc.cuda(loss_test.device.index)
+
+        # in DP mode (default) make sure if result is scalar, there's another dim in the beginning
+        if self.trainer.use_dp:
+            loss_test = loss_test.unsqueeze(0)
+            test_acc = test_acc.unsqueeze(0)
+
+        # alternate possible outputs to test
+        if batch_i % 1 == 0:
+            output = OrderedDict({
+                'test_loss': loss_test,
+                'test_acc': test_acc,
+            })
+            return output
+        if batch_i % 2 == 0:
+            return test_acc
+
+        if batch_i % 3 == 0:
+            output = OrderedDict({
+                'test_loss': loss_test,
+                'test_acc': test_acc,
+                'test_dic': {'test_loss_a': loss_test}
+            })
+            return output
+        if batch_i % 5 == 0:
+            output = OrderedDict({
+                f'test_loss_{dataloader_i}': loss_test,
+                f'test_acc_{dataloader_i}': test_acc,
+            })
+            return output
+
+    def test_end(self, outputs):
+        """
+        Called at the end of validation to aggregate outputs
+        :param outputs: list of individual outputs of each validation step
+        :return:
+        """
+        # if returned a scalar from test_step, outputs is a list of tensor scalars
+        # we return just the average in this case (if we want)
+        # return torch.stack(outputs).mean()
+        test_loss_mean = 0
+        test_acc_mean = 0
+        for output in outputs:
+            test_loss = output['test_loss']
+
+            # reduce manually when using dp
+            if self.trainer.use_dp:
+                test_loss = torch.mean(test_loss)
+            test_loss_mean += test_loss
+
+            # reduce manually when using dp
+            test_acc = output['test_acc']
+            if self.trainer.use_dp:
+                test_acc = torch.mean(test_acc)
+
+            test_acc_mean += test_acc
+
+        test_loss_mean /= len(outputs)
+        test_acc_mean /= len(outputs)
+
+        tqdm_dic = {'test_loss': test_loss_mean.item(), 'test_acc': test_acc_mean.item()}
+        return tqdm_dic
+
     def on_tng_metrics(self, logs):
         logs['some_tensor_to_test'] = torch.rand(1)
 
@@ -235,6 +327,8 @@ class LightningTestModel(LightningModule):
 
     @data_loader
     def test_dataloader(self):
+        if self.use_two_test_sets:
+            return [self.__dataloader(train=False), self.__dataloader(train=False)]
         return self.__dataloader(train=False)
 
     @staticmethod

tests/debug.py

@@ -1,5 +1,6 @@
 from pytorch_lightning import Trainer
 from examples import LightningTemplateModel
+from pytorch_lightning.testing import LightningTestModel, NoValEndTestModel, NoValModel
 from argparse import Namespace
 from test_tube import Experiment
 from pytorch_lightning.callbacks import ModelCheckpoint
@@ -12,6 +13,7 @@ from torch.nn import functional as F
 from torch.utils.data import DataLoader
 from torchvision.datasets import MNIST
 import numpy as np
+import pdb
 
 
 class CoolModel(pl.LightningModule):
@@ -73,10 +75,11 @@ def get_model():
     return model, hparams
 
 
-def get_exp(debug=True):
+def get_exp(debug=True, version=None):
     # set up exp object without actually saving logs
     root_dir = os.path.dirname(os.path.realpath(__file__))
-    exp = Experiment(debug=debug, save_dir=root_dir, name='tests_tt_dir')
+    save_dir = os.path.join(root_dir, 'save_dir')
+    exp = Experiment(debug=debug, save_dir=save_dir, name='tests_tt_dir', version=version)
     return exp
 
 
@@ -99,7 +102,7 @@ def clear_save_dir():
         shutil.rmtree(save_dir)
 
 
-def load_model(exp, save_dir):
+def load_model(exp, save_dir, on_gpu, map_location=None, module_class=LightningTemplateModel):
 
     # load trained model
     tags_path = exp.get_data_path(exp.name, exp.version)
@@ -108,8 +111,10 @@ def load_model(exp, save_dir):
     checkpoints = [x for x in os.listdir(save_dir) if '.ckpt' in x]
     weights_dir = os.path.join(save_dir, checkpoints[0])
 
-    trained_model = LightningTemplateModel.load_from_metrics(weights_path=weights_dir,
+    trained_model = module_class.load_from_metrics(weights_path=weights_dir,
-                                                             tags_csv=tags_path, on_gpu=True)
+                                                   tags_csv=tags_path,
+                                                   on_gpu=on_gpu,
+                                                   map_location=map_location)
 
     assert trained_model is not None, 'loading model failed'
 
@@ -131,12 +136,10 @@ def run_prediction(dataloader, trained_model):
     val_acc = torch.sum(y == labels_hat).item() / (len(y) * 1.0)
     val_acc = torch.tensor(val_acc)
     val_acc = val_acc.item()
-
-    print(val_acc)
-
     assert val_acc > 0.70, 'this model is expected to get > 0.7 in test set (it got %f)' % val_acc
 
 
+# ------------------------------------------------------------------------
 def run_gpu_model_test(trainer_options, model, hparams, on_gpu=True):
     save_dir = init_save_dir()
 
@@ -162,7 +165,7 @@ def run_gpu_model_test(trainer_options, model, hparams, on_gpu=True):
     # test model loading
     pretrained_model = load_model(exp, save_dir, on_gpu)
 
-    # test model preds
+    # test new model accuracy
     run_prediction(model.test_dataloader, pretrained_model)
 
     if trainer.use_ddp:
@@ -177,19 +180,79 @@ def run_gpu_model_test(trainer_options, model, hparams, on_gpu=True):
     clear_save_dir()
 
 
+def assert_ok_val_acc(trainer):
+    # this model should get 0.80+ acc
+    acc = trainer.tng_tqdm_dic['val_acc']
+    assert acc > 0.50, f'model failed to get expected 0.50 validation accuracy. Got: {acc}'
+
+
+def assert_ok_test_acc(trainer):
+    # this model should get 0.80+ acc
+    acc = trainer.tng_tqdm_dic['test_acc']
+    assert acc > 0.50, f'model failed to get expected 0.50 validation accuracy. Got: {acc}'
+
+
+def get_hparams(continue_training=False, hpc_exp_number=0):
+    root_dir = os.path.dirname(os.path.realpath(__file__))
+
+    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
+
+
 def main():
+    """Verify test() on fitted model"""
+    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()
+
+    # exp file to get weights
+    checkpoint = ModelCheckpoint(save_dir)
 
-    os.environ['MASTER_PORT'] = str(np.random.randint(12000, 19000, 1)[0])
-    model, hparams = get_model()
     trainer_options = dict(
+        show_progress_bar=False,
         max_nb_epochs=1,
         train_percent_check=0.4,
         val_percent_check=0.2,
+        checkpoint_callback=checkpoint,
+        experiment=exp,
         gpus=[0, 1],
         distributed_backend='ddp'
     )
 
-    run_gpu_model_test(trainer_options, model, hparams)
+    # fit model
+    trainer = Trainer(**trainer_options)
+    result = trainer.fit(model)
+
+    # correct result and ok accuracy
+    assert result == 1, 'training failed to complete'
+    pretrained_model = load_model(exp, save_dir, on_gpu=True, module_class=LightningTestModel)
+
+    new_trainer = Trainer(**trainer_options)
+    new_trainer.test(pretrained_model)
+
+    # test we have good test accuracy
+    assert_ok_test_acc(new_trainer)
+    # clear_save_dir()
 
 
 if __name__ == '__main__':

tests/test_models.py

@@ -30,6 +30,177 @@ np.random.seed(SEED)
 # ------------------------------------------------------------------------
 # TESTS
 # ------------------------------------------------------------------------
+def test_running_test_pretrained_model_ddp():
+    """Verify test() on pretrained model"""
+    if not can_run_gpu_test():
+        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()
+
+    # exp file to get weights
+    checkpoint = ModelCheckpoint(save_dir)
+
+    trainer_options = dict(
+        show_progress_bar=False,
+        max_nb_epochs=1,
+        train_percent_check=0.4,
+        val_percent_check=0.2,
+        checkpoint_callback=checkpoint,
+        experiment=exp,
+        gpus=[0, 1],
+        distributed_backend='ddp'
+    )
+
+    # fit model
+    trainer = Trainer(**trainer_options)
+    result = trainer.fit(model)
+
+    # correct result and ok accuracy
+    assert result == 1, 'training failed to complete'
+    pretrained_model = load_model(exp, save_dir, on_gpu=True, module_class=LightningTestModel)
+
+    # run test set
+    new_trainer = Trainer(**trainer_options)
+    new_trainer.test(pretrained_model)
+
+    run_prediction(model.test_dataloader, pretrained_model)
+
+    # test we have good test accuracy
+    clear_save_dir()
+
+
+def test_running_test_after_fitting():
+    """Verify test() on fitted model"""
+    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()
+
+    # exp file to get weights
+    checkpoint = ModelCheckpoint(save_dir)
+
+    trainer_options = dict(
+        show_progress_bar=False,
+        max_nb_epochs=1,
+        train_percent_check=0.4,
+        val_percent_check=0.2,
+        test_percent_check=0.2,
+        checkpoint_callback=checkpoint,
+        experiment=exp
+    )
+
+    # fit model
+    trainer = Trainer(**trainer_options)
+    result = trainer.fit(model)
+
+    assert result == 1, 'training failed to complete'
+
+    trainer.test()
+
+    # test we have good test accuracy
+    assert_ok_test_acc(trainer)
+
+    clear_save_dir()
+
+
+def test_running_test_pretrained_model():
+    """Verify test() on pretrained model"""
+    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()
+
+    # exp file to get weights
+    checkpoint = ModelCheckpoint(save_dir)
+
+    trainer_options = dict(
+        show_progress_bar=False,
+        max_nb_epochs=1,
+        train_percent_check=0.4,
+        val_percent_check=0.2,
+        checkpoint_callback=checkpoint,
+        experiment=exp
+    )
+
+    # fit model
+    trainer = Trainer(**trainer_options)
+    result = trainer.fit(model)
+
+    # correct result and ok accuracy
+    assert result == 1, 'training failed to complete'
+    pretrained_model = load_model(exp, save_dir, on_gpu=False, module_class=LightningTestModel)
+
+    new_trainer = Trainer(**trainer_options)
+    new_trainer.test(pretrained_model)
+
+    # test we have good test accuracy
+    assert_ok_test_acc(new_trainer)
+    clear_save_dir()
+
+
+def test_running_test_pretrained_model_dp():
+    """Verify test() on pretrained model"""
+    if not can_run_gpu_test():
+        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()
+
+    # exp file to get weights
+    checkpoint = ModelCheckpoint(save_dir)
+
+    trainer_options = dict(
+        show_progress_bar=False,
+        max_nb_epochs=1,
+        train_percent_check=0.4,
+        val_percent_check=0.2,
+        checkpoint_callback=checkpoint,
+        experiment=exp,
+        gpus=[0, 1],
+        distributed_backend='dp'
+    )
+
+    # fit model
+    trainer = Trainer(**trainer_options)
+    result = trainer.fit(model)
+
+    # correct result and ok accuracy
+    assert result == 1, 'training failed to complete'
+    pretrained_model = load_model(exp, save_dir, on_gpu=True, module_class=LightningTestModel)
+
+    new_trainer = Trainer(**trainer_options)
+    new_trainer.test(pretrained_model)
+
+    # test we have good test accuracy
+    assert_ok_test_acc(new_trainer)
+    clear_save_dir()
+
+
 def test_gradient_accumulation_scheduling():
     """
     Test grad accumulation by the freq of optimizer updates
@@ -107,13 +278,7 @@ def test_multi_gpu_model_ddp():
     Make sure DDP works
     :return:
     """
-    if not torch.cuda.is_available():
+    if not can_run_gpu_test():
-        warnings.warn('test_multi_gpu_model_ddp cannot run.'
-                      ' Rerun on a GPU node to run this test')
-        return
-    if not torch.cuda.device_count() > 1:
-        warnings.warn('test_multi_gpu_model_ddp cannot run.'
-                      ' Rerun on a node with 2+ GPUs to run this test')
         return
 
     os.environ['MASTER_PORT'] = str(np.random.randint(12000, 19000, 1)[0])
@@ -161,13 +326,7 @@ def test_optimizer_return_options():
 
 
 def test_single_gpu_batch_parse():
-    if not torch.cuda.is_available():
+    if not can_run_gpu_test():
-        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()
@@ -261,7 +420,7 @@ def test_no_val_module():
     trainer = Trainer(**trainer_options)
     result = trainer.fit(model)
 
-    # traning complete
+    # training complete
     assert result == 1, 'amp + ddp model failed to complete'
 
     # save model
@@ -449,13 +608,7 @@ def test_amp_gpu_ddp():
     Make sure DDP + AMP work
     :return:
     """
-    if not torch.cuda.is_available():
+    if not can_run_gpu_test():
-        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])
@@ -667,13 +820,7 @@ def test_amp_gpu_ddp_slurm_managed():
     Make sure DDP + AMP work
     :return:
     """
-    if not torch.cuda.is_available():
+    if not can_run_gpu_test():
-        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
 
     # simulate setting slurm flags
@@ -832,14 +979,9 @@ def test_multi_gpu_model_dp():
     Make sure DP works
     :return:
     """
-    if not torch.cuda.is_available():
+    if not can_run_gpu_test():
-        warnings.warn('test_multi_gpu_model_dp cannot run.'
-                      ' Rerun on a GPU node to run this test')
-        return
-    if not torch.cuda.device_count() > 1:
-        warnings.warn('test_multi_gpu_model_dp cannot run.'
-                      ' Rerun on a node with 2+ GPUs to run this test')
         return
+
     model, hparams = get_model()
     trainer_options = dict(
         show_progress_bar=False,
@@ -860,14 +1002,9 @@ def test_amp_gpu_dp():
     Make sure DP + AMP work
     :return:
     """
-    if not torch.cuda.is_available():
+    if not can_run_gpu_test():
-        warnings.warn('test_amp_gpu_dp cannot run.'
-                      ' Rerun on a GPU node to run this test')
-        return
-    if not torch.cuda.device_count() > 1:
-        warnings.warn('test_amp_gpu_dp cannot run.'
-                      ' Rerun on a node with 2+ GPUs to run this test')
         return
+
     model, hparams = get_model()
     trainer_options = dict(
         max_nb_epochs=1,
@@ -884,11 +1021,7 @@ def test_ddp_sampler_error():
     Make sure DDP + AMP work
     :return:
     """
-    if not torch.cuda.is_available():
+    if not can_run_gpu_test():
-        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])
@@ -922,7 +1055,34 @@ def test_multiple_val_dataloader():
     hparams = get_hparams()
     model = LightningTestModel(hparams)
 
-    save_dir = init_save_dir()
+    # exp file to get meta
+    trainer_options = dict(
+        max_nb_epochs=1,
+        val_percent_check=0.1,
+        train_percent_check=1.0,
+    )
+
+    # fit model
+    trainer = Trainer(**trainer_options)
+    result = trainer.fit(model)
+
+    # verify tng completed
+    assert result == 1
+
+    # verify there are 2 val loaders
+    assert len(trainer.val_dataloader) == 2, 'Multiple val_dataloaders not initiated properly'
+
+    # make sure predictions are good for each val set
+    [run_prediction(dataloader, trainer.model) for dataloader in trainer.val_dataloader]
+
+
+def test_multiple_test_dataloader():
+    """
+    Verify multiple test_dataloader
+    :return:
+    """
+    hparams = get_hparams()
+    model = LightningTestModel(hparams, use_two_test_sets=True)
 
     # exp file to get meta
     trainer_options = dict(
@@ -939,10 +1099,10 @@ def test_multiple_val_dataloader():
     assert result == 1
 
     # verify there are 2 val loaders
-    assert len(trainer.val_dataloader) == 2, 'Multiple val_dataloaders not initiated properly'
+    assert len(trainer.test_dataloader) == 2, 'Multiple test_dataloaders not initiated properly'
 
-    # make sure predictions are good for each val set
+    # make sure predictions are good for each test set
-    [run_prediction(dataloader, trainer.model) for dataloader in trainer.val_dataloader]
+    [run_prediction(dataloader, trainer.model) for dataloader in trainer.test_dataloader]
 
 
 # ------------------------------------------------------------------------
@@ -973,7 +1133,7 @@ def run_gpu_model_test(trainer_options, model, hparams, on_gpu=True):
     # test model loading
     pretrained_model = load_model(exp, save_dir, on_gpu)
 
-    # test model preds
+    # test new model accuracy
     run_prediction(model.test_dataloader, pretrained_model)
 
     if trainer.use_ddp:
@@ -1024,7 +1184,8 @@ def get_model(use_test_model=False):
 def get_exp(debug=True, version=None):
     # set up exp object without actually saving logs
     root_dir = os.path.dirname(os.path.realpath(__file__))
-    exp = Experiment(debug=debug, save_dir=root_dir, name='tests_tt_dir', version=version)
+    save_dir = os.path.join(root_dir, 'save_dir')
+    exp = Experiment(debug=debug, save_dir=save_dir, name='tests_tt_dir', version=version)
     return exp
 
 
@@ -1033,7 +1194,8 @@ def init_save_dir():
     save_dir = os.path.join(root_dir, 'save_dir')
 
     if os.path.exists(save_dir):
-        shutil.rmtree(save_dir)
+        n = np.random.randint(0, 10000000, 1)[0]
+        shutil.move(save_dir, save_dir + f'_{n}')
 
     os.makedirs(save_dir, exist_ok=True)
 
@@ -1044,10 +1206,11 @@ def clear_save_dir():
     root_dir = os.path.dirname(os.path.realpath(__file__))
     save_dir = os.path.join(root_dir, 'save_dir')
     if os.path.exists(save_dir):
-        shutil.rmtree(save_dir)
+        n = np.random.randint(0, 10000000, 1)[0]
+        shutil.move(save_dir, save_dir + f'_{n}')
 
 
-def load_model(exp, save_dir, on_gpu, map_location=None):
+def load_model(exp, save_dir, on_gpu, map_location=None, module_class=LightningTemplateModel):
 
     # load trained model
     tags_path = exp.get_data_path(exp.name, exp.version)
@@ -1056,10 +1219,10 @@ def load_model(exp, save_dir, on_gpu, map_location=None):
     checkpoints = [x for x in os.listdir(save_dir) if '.ckpt' in x]
     weights_dir = os.path.join(save_dir, checkpoints[0])
 
-    trained_model = LightningTemplateModel.load_from_metrics(weights_path=weights_dir,
+    trained_model = module_class.load_from_metrics(weights_path=weights_dir,
-                                                             tags_csv=tags_path,
+                                                   tags_csv=tags_path,
-                                                             on_gpu=on_gpu,
+                                                   on_gpu=on_gpu,
-                                                             map_location=map_location)
+                                                   map_location=map_location)
 
     assert trained_model is not None, 'loading model failed'
 
@@ -1078,19 +1241,35 @@ def run_prediction(dataloader, trained_model):
 
     # acc
     labels_hat = torch.argmax(y_hat, dim=1)
-    val_acc = torch.sum(y == labels_hat).item() / (len(y) * 1.0)
+    acc = torch.sum(y == labels_hat).item() / (len(y) * 1.0)
-    val_acc = torch.tensor(val_acc)
+    acc = torch.tensor(acc)
-    val_acc = val_acc.item()
+    acc = acc.item()
 
-    print(val_acc)
+    assert acc > 0.50, f'this model is expected to get > 0.50 in test set (it got {acc})'
-
-    assert val_acc > 0.50, 'this model is expected to get > 0.50 in test set (it got %f)' % val_acc
 
 
-def assert_ok_acc(trainer):
+def assert_ok_val_acc(trainer):
     # this model should get 0.80+ acc
     acc = trainer.tng_tqdm_dic['val_acc']
-    assert acc > 0.50, 'model failed to get expected 0.50 validation accuracy. Got: %f' % acc
+    assert acc > 0.50, f'model failed to get expected 0.50 validation accuracy. Got: {acc}'
+
+
+def assert_ok_test_acc(trainer):
+    # this model should get 0.80+ acc
+    acc = trainer.tng_tqdm_dic['test_acc']
+    assert acc > 0.50, f'model failed to get expected 0.50 validation accuracy. Got: {acc}'
+
+
+def can_run_gpu_test():
+    if not torch.cuda.is_available():
+        warnings.warn('test_multi_gpu_model_ddp cannot run.'
+                      ' Rerun on a GPU node to run this test')
+        return False
+    if not torch.cuda.device_count() > 1:
+        warnings.warn('test_multi_gpu_model_ddp cannot run.'
+                      ' Rerun on a node with 2+ GPUs to run this test')
+        return False
+    return True
 
 
 if __name__ == '__main__':