Repository: Lightning-AI/pytorch-lightning

Issue:
Pass output from all val dataloaders to `validation_end`
**Is your feature request related to a problem? Please describe.**
I'm working on a retrieval problem. To compute validation scores, I'd like to compute feature vectors for a query set and a catalog set, then measure information retrieval metrics (precision, recall, etc) of the query set against the catalog set.

To do this, I'd like to return the two datasets from `val_dataloader` (already possible), then use `validation_step` to accumulate all the feature vectors for each dataset (already possible). Once I have all the feature vectors, I need some way to bring the vectors from the two datasets together so that I can compute retrieval metrics. Right now, `validation_end` processes the outputs of multiple validation dataloaders separately.

**Describe the solution you'd like**
Change the behavior of `validation_end` in cases where there are multiple validation dataloaders. I propose that it only be called once, but that it be passed a list of outputs dicts, one from each dataloader. Using it would look something like this:

```python
def validation_end(self, outputs):
    query_output, catalog_output = outputs
    query_vectors = torch.cat([o['vectors'] for o in query_output])
    query_labels = torch.cat([o['labels'] for o in query_output])
    catalog_vectors = torch.cat([o['vectors'] for o in catalog_output])
    catalog_labels = torch.cat([o['labels'] for o in catalog_output])
    return retrieval_metrics(
        query_vectors,
        query_labels,
        catalog_vectors,
        catalog_labels
    )
```

**Describe alternatives you've considered**
* Breaking outside of lightning entirely to do validation
* Concatenating query and catalog datasets, and somehow adding a flag to indicate which dataset a given example came from




Discussion:
good point. do you want to take a stab at implementing? otherwise one of us can look at this end of week
Sure, I can take a look at it over the weekend.
also needs to happen for test_end

Resolving patch:
diff --git a/pytorch_lightning/models/trainer.py b/pytorch_lightning/models/trainer.py
--- a/pytorch_lightning/models/trainer.py
+++ b/pytorch_lightning/models/trainer.py
@@ -423,11 +423,11 @@ def __evaluation_forward(self, model, data_batch, batch_i, dataloader_i, test=Fa
 
         return output
 
-    def evaluate(self, model, dataloader, max_batches, dataloader_i, test=False):
+    def evaluate(self, model, dataloaders, max_batches, test=False):
         """
         Run evaluation code
         :param model: PT model
-        :param dataloader: PT dataloader
+        :param dataloaders: list of PT dataloaders
         :param max_batches: Scalar
         :param dataloader_i:
         :param test: boolean
@@ -444,32 +444,37 @@ def evaluate(self, model, dataloader, max_batches, dataloader_i, test=False):
         outputs = []
 
         # run training
-        for batch_i, data_batch in enumerate(dataloader):
+        for dataloader_i, dl in enumerate(dataloaders):
+            dl_outputs = []
+            for batch_i, data_batch in enumerate(dl):
 
-            if data_batch is None:  # pragma: no cover
-                continue
+                if data_batch is None:  # pragma: no cover
+                    continue
 
-            # stop short when on fast_dev_run (sets max_batch=1)
-            if batch_i >= max_batches:
-                break
+                # stop short when on fast_dev_run (sets max_batch=1)
+                if batch_i >= max_batches:
+                    break
 
-            # -----------------
-            # RUN EVALUATION STEP
-            # -----------------
-            output = self.__evaluation_forward(model, data_batch, batch_i, dataloader_i,
-                                               test)
+                # -----------------
+                # RUN EVALUATION STEP
+                # -----------------
+                output = self.__evaluation_forward(model, data_batch, batch_i, dataloader_i,
+                                                   test)
 
-            # track outputs for collation
-            outputs.append(output)
+                # track outputs for collation
+                dl_outputs.append(output)
 
-            # batch done
-            if self.show_progress_bar:
-                self.progress_bar.update(1)
+                # batch done
+                if self.show_progress_bar:
+                    self.progress_bar.update(1)
+            outputs.append(dl_outputs)
 
         eval_results = {}
 
         # give model a chance to do something with the outputs (and method defined)
         model = self.__get_model()
+        if len(dataloaders) == 1:
+            outputs = outputs[0]
         if test and self.__is_overriden('test_end'):
             eval_results = model.test_end(outputs)
         elif self.__is_overriden('validation_end'):
@@ -837,13 +842,11 @@ def __run_pretrain_routine(self, model):
         # 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):
-
-                # reset progress_bar limit for sanity check
-                if self.show_progress_bar:
-                    self.progress_bar.reset(self.nb_sanity_val_steps)
+            # reset progress_bar limit for sanity check
+            if self.show_progress_bar:
+                self.progress_bar.reset(self.nb_sanity_val_steps)
 
-                self.evaluate(model, dataloader, self.nb_sanity_val_steps, ds_i, self.testing)
+            self.evaluate(model, self.val_dataloader, self.nb_sanity_val_steps, self.testing)
 
         # ---------------------------
         # CORE TRAINING LOOP
@@ -1203,17 +1206,15 @@ def __run_evaluation(self, test=False):
             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)
+            eval_out_metrics = self.evaluate(self.model,
+                                             dataloaders,
+                                             max_batches,
+                                             test)
 
-                self.__add_tqdm_metrics(eval_out_metrics)
+            self.__add_tqdm_metrics(eval_out_metrics)
 
-                # hook
-                model.on_post_performance_check()
+            # hook
+            model.on_post_performance_check()
 
             if self.show_progress_bar:
                 # add model specific metrics


Test patch:
diff --git a/pytorch_lightning/testing/lm_test_module_mixins.py b/pytorch_lightning/testing/lm_test_module_mixins.py
--- a/pytorch_lightning/testing/lm_test_module_mixins.py
+++ b/pytorch_lightning/testing/lm_test_module_mixins.py
@@ -185,23 +185,26 @@ def validation_end(self, outputs):
         # return torch.stack(outputs).mean()
         val_loss_mean = 0
         val_acc_mean = 0
-        for output in outputs:
-            val_loss = output['val_loss']
+        i = 0
+        for dl_output in outputs:
+            for output in dl_output:
+                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
+                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)
+                # 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_acc_mean += val_acc
+                i += 1
 
-        val_loss_mean /= len(outputs)
-        val_acc_mean /= len(outputs)
+        val_loss_mean /= i
+        val_acc_mean /= i
 
         tqdm_dic = {'val_loss': val_loss_mean.item(), 'val_acc': val_acc_mean.item()}
         return tqdm_dic
@@ -359,23 +362,26 @@ def test_end(self, outputs):
         # 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)
+        i = 0
+        for dl_output in outputs:
+            for output in dl_output:
+                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
+                i += 1
+
+        test_loss_mean /= i
+        test_acc_mean /= i
 
         tqdm_dic = {'test_loss': test_loss_mean.item(), 'test_acc': test_acc_mean.item()}
         return tqdm_dic
diff --git a/tests/test_models.py b/tests/test_models.py
--- a/tests/test_models.py
+++ b/tests/test_models.py
@@ -1162,15 +1162,15 @@ class CurrentTestModel(
     trainer = Trainer(**trainer_options)
     result = trainer.fit(model)
 
-    # verify tng completed
-    assert result == 1
-
     # verify there are 2 val loaders
     assert len(trainer.test_dataloader) == 2, 'Multiple test_dataloaders not initiated properly'
 
     # make sure predictions are good for each test set
     [run_prediction(dataloader, trainer.model) for dataloader in trainer.test_dataloader]
 
+    # run the test method
+    trainer.test()
+
 
 # ------------------------------------------------------------------------
 # UTILS