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
2026-09-09 11:32:07 +08:00 · 2019-08-30 18:56:09 -04:00
parent 73cf47112e
commit f51b45933b
8 changed files with 707 additions and 147 deletions
@@ -16,7 +16,9 @@ Otherwise, to Define a Lightning Module, implement the following methods:
 **Optional**:   

 - [validation_step](RequiredTrainerInterface.md#validation_step)    
- [validation_end](RequiredTrainerInterface.md#validation_end)    
+- [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 

@@ -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
-        self.nb_test_batches = int(self.nb_test_batches * self.test_percent_check)
+        if self.test_dataloader is not None:
+            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:
-            output = model(*args)
-        elif self.single_gpu:
-            # put inputs on gpu manually
+            return output
+
+        # CPU, single GPU
+        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
-            output = model.validation_step(*args)
-
+        if test:
+            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 = {}
-        if self.__is_overriden('validation_end'):
-            if self.data_parallel:
-                val_results = model.module.validation_end(outputs)
-            else:
-                val_results = model.validation_end(outputs)
+        model = self.__get_model()
+        if test and self.__is_overriden('test_end'):
+            eval_results = model.test_end(outputs)
+        elif self.__is_overriden('validation_end'):
+            eval_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
-        ref_model.trainer = self
-
+        # link up experiment object
        if self.experiment is not None:
            ref_model.experiment = self.experiment

-        # save exp to get started
-        if self.proc_rank == 0 and self.experiment is not None:
-            self.experiment.save()
+            # save exp to get started
+            if self.proc_rank == 0:
+                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):
-        # decide if can check epochs
-        can_check_epoch = (self.current_epoch + 1) % self.check_val_every_n_epoch == 0
-        if self.fast_dev_run:
-            print('skipping to check performance bc of --fast_dev_run')
-        elif not can_check_epoch:
-            return
+    def __run_evaluation(self, test=False):
+        # when testing make sure user defined a test step
+        can_run_test_step = False
+        if test:
+            can_run_test_step = self.__is_overriden('test_step') and self.__is_overriden('test_end')
+            if not can_run_test_step:
+                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.on_pre_performance_check()
+            model = self.__get_model()
+            model.on_pre_performance_check()

-            # use val_percent_check set on end of epoch
-            # use a small portion otherwise
-            max_batches = self.nb_val_batches if not self.fast_dev_run else 1
-            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)
+            # select dataloaders
+            dataloaders = self.val_dataloader
+            max_batches = self.nb_val_batches

-            # hook
-            if self.__is_function_implemented('on_post_performance_check'):
-                model = self.__get_model()
+            # calculate max batches to use
+            if test:
+                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)
@@ -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)
                # ---------------
@@ -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
@@ -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)
@@ -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_acc_mean += output['val_acc']
+            val_loss = output['val_loss']
+
+            # 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
@@ -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,
-                                                             tags_csv=tags_path, on_gpu=True)
+    trained_model = module_class.load_from_metrics(weights_path=weights_dir,
+                                                   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__':
@@ -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():
-        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')
+    if not can_run_gpu_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():
-        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')
+    if not can_run_gpu_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():
-        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')
+    if not can_run_gpu_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():
-        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')
+    if not can_run_gpu_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():
-        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')
+    if not can_run_gpu_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():
-        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')
+    if not can_run_gpu_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():
-        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')
+    if not can_run_gpu_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
-    [run_prediction(dataloader, trainer.model) for dataloader in trainer.val_dataloader]
+    # make sure predictions are good for each test set
+    [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,
-                                                             tags_csv=tags_path,
-                                                             on_gpu=on_gpu,
-                                                             map_location=map_location)
+    trained_model = module_class.load_from_metrics(weights_path=weights_dir,
+                                                   tags_csv=tags_path,
+                                                   on_gpu=on_gpu,
+                                                   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)
-    val_acc = torch.tensor(val_acc)
-    val_acc = val_acc.item()
+    acc = torch.sum(y == labels_hat).item() / (len(y) * 1.0)
+    acc = torch.tensor(acc)
+    acc = acc.item()

-    print(val_acc)
-
-    assert val_acc > 0.50, 'this model is expected to get > 0.50 in test set (it got %f)' % val_acc
+    assert acc > 0.50, f'this model is expected to get > 0.50 in test set (it got {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__':