Fix logging to loggers with multiple eval dataloaders (#12454)

Author: rohitgr7

CHANGELOG.md

@@ -924,6 +924,9 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Fixed initializing optimizers unnecessarily in `DDPFullyShardedStrategy` ([#12267](https://github.com/PyTorchLightning/pytorch-lightning/pull/12267))
 
 
+- Fixed logging to loggers with multiple eval dataloaders ([#12454](https://github.com/PyTorchLightning/pytorch-lightning/pull/12454))
+
+
 ## [1.5.10] - 2022-02-08
 
 ### Fixed

pytorch_lightning/loops/dataloader/evaluation_loop.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 import os
 import shutil
-from collections import OrderedDict
+from collections import ChainMap, OrderedDict
 from functools import partial
 from typing import Any, IO, Iterable, List, Optional, Sequence, Type, Union
 
@@ -181,11 +181,13 @@ 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
         epoch_end_logged_outputs = self.trainer._logger_connector.update_eval_epoch_metrics()
+        all_logged_outputs = dict(ChainMap(*logged_outputs))  # list[dict] -> dict
+        all_logged_outputs.update(epoch_end_logged_outputs)
         for dl_outputs in logged_outputs:
             dl_outputs.update(epoch_end_logged_outputs)
 
         # log metrics
-        self.trainer._logger_connector.log_eval_end_metrics()
+        self.trainer._logger_connector.log_eval_end_metrics(all_logged_outputs)
 
         # hook
         self._on_evaluation_end()

pytorch_lightning/trainer/connectors/logger_connector/logger_connector.py

@@ -174,13 +174,13 @@ def update_eval_epoch_metrics(self) -> _OUT_DICT:
         self._logged_metrics.update(metrics["log"])
         return metrics["log"]
 
-    def log_eval_end_metrics(self) -> None:
+    def log_eval_end_metrics(self, metrics: _OUT_DICT) -> None:
         assert self._epoch_end_reached
         if self.trainer.sanity_checking:
             return
 
         # log all the metrics as a single dict
-        self.log_metrics(self.metrics["log"])
+        self.log_metrics(metrics)
 
     """
     Train metric updates

pytorch_lightning/utilities/imports.py

@@ -89,6 +89,7 @@ def _compare_version(package: str, op: Callable, version: str, use_base_version:
 
 _IS_WINDOWS = platform.system() == "Windows"
 _IS_INTERACTIVE = hasattr(sys, "ps1")  # https://stackoverflow.com/a/64523765
+_PYTHON_GREATER_EQUAL_3_8_0 = _compare_version("python", operator.ge, "3.8.0")
 _TORCH_GREATER_EQUAL_1_8_1 = _compare_version("torch", operator.ge, "1.8.1")
 _TORCH_GREATER_EQUAL_1_9 = _compare_version("torch", operator.ge, "1.9.0")
 _TORCH_GREATER_EQUAL_1_9_1 = _compare_version("torch", operator.ge, "1.9.1")

tests/loops/test_evaluation_loop.py

@@ -48,7 +48,7 @@ def test_on_evaluation_epoch_end(eval_epoch_end_mock, tmpdir):
 def test_log_epoch_metrics_before_on_evaluation_end(update_eval_epoch_metrics_mock, tmpdir):
     """Test that the epoch metrics are logged before the `on_evaluation_end` hook is fired."""
     order = []
-    update_eval_epoch_metrics_mock.side_effect = lambda: order.append("log_epoch_metrics")
+    update_eval_epoch_metrics_mock.side_effect = lambda _: order.append("log_epoch_metrics")
 
     class LessBoringModel(BoringModel):
         def on_validation_end(self):

tests/trainer/logging_/test_eval_loop_logging.py

@@ -28,6 +28,7 @@
 from pytorch_lightning.loops.dataloader import EvaluationLoop
 from pytorch_lightning.trainer.states import RunningStage
 from pytorch_lightning.utilities.exceptions import MisconfigurationException
+from pytorch_lightning.utilities.imports import _PYTHON_GREATER_EQUAL_3_8_0
 from tests.helpers import BoringModel, RandomDataset
 from tests.helpers.runif import RunIf
 
@@ -573,10 +574,8 @@ def test_step(self, batch, batch_idx):
     assert mock_log_metrics.mock_calls[0] == call({"hp_metric": -1}, 0)
 
     def get_metrics_at_idx(idx):
-        mock_calls = list(mock_log_metrics.mock_calls)
-        if isinstance(mock_calls[idx].kwargs, dict):
-            return mock_calls[idx].kwargs["metrics"]
-        return mock_calls[idx][2]["metrics"]
+        mock_call = mock_log_metrics.mock_calls[idx]
+        return mock_call.kwargs["metrics"] if _PYTHON_GREATER_EQUAL_3_8_0 else mock_call[2]["metrics"]
 
     expected = {"valid_loss_0_step", "valid_loss_2"}
     assert set(get_metrics_at_idx(1)) == expected
@@ -755,7 +754,8 @@ def test_dataloader(self):
     }
 
 
-def test_logging_multi_dataloader_on_epoch_end(tmpdir):
+@mock.patch("pytorch_lightning.loggers.TensorBoardLogger.log_metrics")
+def test_logging_multi_dataloader_on_epoch_end(mock_log_metrics, tmpdir):
     class CustomBoringModel(BoringModel):
         def test_step(self, batch, batch_idx, dataloader_idx):
             self.log("foo", dataloader_idx + 1)
@@ -765,13 +765,21 @@ def test_epoch_end(self, outputs) -> None:
             self.log("foobar", sum(sum(o) for o in outputs))
 
         def test_dataloader(self):
-            return [super().val_dataloader(), super().val_dataloader()]
+            return [super().test_dataloader(), super().test_dataloader()]
 
     model = CustomBoringModel()
-    trainer = Trainer(default_root_dir=tmpdir, fast_dev_run=1)
+    trainer = Trainer(default_root_dir=tmpdir, limit_test_batches=1)
     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}]
+    cb_metrics = set(trainer.callback_metrics)
+    assert cb_metrics == {"foo/dataloader_idx_0", "foo/dataloader_idx_1", "foobar"}
+
+    mock_call = mock_log_metrics.mock_calls[0]
+    logged_metrics = mock_call.kwargs["metrics"] if _PYTHON_GREATER_EQUAL_3_8_0 else mock_call[2]["metrics"]
+    cb_metrics.add("epoch")
+    assert set(logged_metrics) == cb_metrics
 
 
 inputs0 = ([{"log": torch.tensor(5)}, {"no_log": torch.tensor(6)}], RunningStage.TESTING)