Fix evaluation logging on epoch end with multiple dataloaders (#11132)

Author: carmocca

CHANGELOG.md

@@ -317,6 +317,9 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Fixed bug where `Trainer(track_grad_norm=..., logger=False)' would fail ([#11114](https://github.com/PyTorchLightning/pytorch-lightning/pull/11114))
 
 
+- Fixed logging on `{test,validation}_epoch_end` with multiple dataloaders ([#11132](https://github.com/PyTorchLightning/pytorch-lightning/pull/11132))
+
+
 - Fixed double evaluation bug with fault-tolerance enabled where the second call was completely skipped ([#11119](https://github.com/PyTorchLightning/pytorch-lightning/pull/11119))
 
 ## [1.5.6] - 2021-12-15

pytorch_lightning/loops/dataloader/evaluation_loop.py

@@ -147,10 +147,9 @@ def on_run_end(self) -> List[_OUT_DICT]:
 
         logged_outputs, self._logged_outputs = self._logged_outputs, []  # free memory
         # include any logged outputs on epoch_end
-        if self.num_dataloaders < 2:  # TODO: remove this check
-            epoch_end_logged_outputs = self.trainer.logger_connector.update_eval_epoch_metrics()
-            for dl_outputs in logged_outputs:
-                dl_outputs.update(epoch_end_logged_outputs)
+        epoch_end_logged_outputs = self.trainer.logger_connector.update_eval_epoch_metrics()
+        for dl_outputs in logged_outputs:
+            dl_outputs.update(epoch_end_logged_outputs)
 
         # log metrics
         self.trainer.logger_connector.log_eval_end_metrics()

pytorch_lightning/trainer/connectors/logger_connector/logger_connector.py

@@ -157,7 +157,11 @@ def update_eval_epoch_metrics(self) -> _OUT_DICT:
         assert self._epoch_end_reached
         if self.trainer.sanity_checking:
             return {}
-        return self.metrics["callback"]
+        metrics = self.metrics
+        self._progress_bar_metrics.update(metrics["pbar"])
+        self._callback_metrics.update(metrics["callback"])
+        self._logged_metrics.update(metrics["log"])
+        return metrics["callback"]
 
     def log_eval_end_metrics(self) -> None:
         assert self._epoch_end_reached

pytorch_lightning/trainer/connectors/logger_connector/result.py

@@ -529,8 +529,7 @@ def valid_items(self) -> Generator:
         return (
             (k, v)
             for k, v in self.items()
-            if not (isinstance(v, _ResultMetric) and v.has_reset)
-            and self.dataloader_idx in (None, v.meta.dataloader_idx)
+            if not (isinstance(v, _ResultMetric) and v.has_reset) and self.dataloader_idx == v.meta.dataloader_idx
         )
 
     def _forked_name(self, result_metric: _ResultMetric, on_step: bool) -> Tuple[str, str]:

tests/trainer/logging_/test_eval_loop_logging.py

@@ -747,16 +747,17 @@ def test_dataloader(self):
 def test_logging_multi_dataloader_on_epoch_end(tmpdir):
     class CustomBoringModel(BoringModel):
         def test_step(self, batch, batch_idx, dataloader_idx):
-            self.log("foo", 12.0)
+            self.log("foo", dataloader_idx + 1)
+            return dataloader_idx + 1
 
         def test_epoch_end(self, outputs) -> None:
-            self.log("foobar", 23.0)
+            self.log("foobar", sum(sum(o) for o in outputs))
 
         def test_dataloader(self):
-            return [torch.utils.data.DataLoader(RandomDataset(32, 64)) for _ in range(2)]
+            return [super().val_dataloader(), super().val_dataloader()]
 
     model = CustomBoringModel()
     trainer = Trainer(default_root_dir=tmpdir, fast_dev_run=1)
-    logged_results = trainer.test(model)
-    # TODO: what's logged in `test_epoch_end` should be included in the results of each dataloader
-    assert logged_results == [{"foo/dataloader_idx_0": 12.0}, {"foo/dataloader_idx_1": 12.0}]
+    results = trainer.test(model)
+    # what's logged in `test_epoch_end` gets included in the results of each dataloader
+    assert results == [{"foo/dataloader_idx_0": 1, "foobar": 3}, {"foo/dataloader_idx_1": 2, "foobar": 3}]