Remove support for logging multiple metrics together (#16389)

Author: carmocca

docs/source-pytorch/extensions/logging.rst

@@ -115,11 +115,6 @@ methods to log from anywhere in a :doc:`LightningModule <../common/lightning_mod
         self.log("my_metric", x)
 
 
-    # or a dict to get multiple metrics on the same plot if the logger supports it
-    def training_step(self, batch, batch_idx):
-        self.log("performance", {"acc": acc, "recall": recall})
-
-
     # or a dict to log all metrics at once with individual plots
     def training_step(self, batch, batch_idx):
         self.log_dict({"acc": acc, "recall": recall})

docs/source-pytorch/visualize/logging_intermediate.rst

@@ -35,17 +35,6 @@ then access the logger's API directly
 
 ----
 
-****************************************
-Track multiple metrics in the same chart
-****************************************
-If your logger supports plotting multiple metrics on the same chart, pass in a dictionary to *self.log*.
-
-.. code-block:: python
-
-    self.log("performance", {"acc": acc, "recall": recall})
-
-----
-
 *********************
 Track hyperparameters
 *********************

src/pytorch_lightning/CHANGELOG.md

@@ -98,6 +98,8 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 - Removed the automatic addition of a moving average of the `training_step` loss in the progress bar. Use `self.log("loss", ..., prog_bar=True)` instead. ([#16192](https://github.com/Lightning-AI/lightning/issues/16192))
 
+- Removed support for passing a dictionary value to `self.log()` ([#16389](https://github.com/Lightning-AI/lightning/pull/16389))
+
 - Tuner removal
   * Removed the deprecated `trainer.tuning` property ([#16379](https://github.com/Lightning-AI/lightning/pull/16379))
   * Removed the deprecated `TrainerFn.TUNING` and `RunningStage.TUNING` enums ([#16379](https://github.com/Lightning-AI/lightning/pull/16379))

src/pytorch_lightning/core/module.py

@@ -53,7 +53,7 @@
 from pytorch_lightning.utilities.rank_zero import rank_zero_debug, rank_zero_warn, WarningCache
 from pytorch_lightning.utilities.signature_utils import is_param_in_hook_signature
 from pytorch_lightning.utilities.types import (
-    _METRIC_COLLECTION,
+    _METRIC,
     EPOCH_OUTPUT,
     LRSchedulerPLType,
     LRSchedulerTypeUnion,
@@ -337,7 +337,7 @@ def forward(self, x):
     def log(
         self,
         name: str,
-        value: _METRIC_COLLECTION,
+        value: _METRIC,
         prog_bar: bool = False,
         logger: Optional[bool] = None,
         on_step: Optional[bool] = None,
@@ -361,7 +361,7 @@ def log(
 
         Args:
             name: key to log.
-            value: value to log. Can be a ``float``, ``Tensor``, ``Metric``, or a dictionary of the former.
+            value: value to log. Can be a ``float``, ``Tensor``, or a ``Metric``.
             prog_bar: if ``True`` logs to the progress bar.
             logger: if ``True`` logs to the logger.
             on_step: if ``True`` logs at this step. The default value is determined by the hook.
@@ -390,7 +390,7 @@ def log(
         # check for invalid values
         apply_to_collection(value, dict, self.__check_not_nested, name)
         apply_to_collection(
-            value, object, self.__check_allowed, name, value, wrong_dtype=(numbers.Number, Metric, Tensor, dict)
+            value, object, self.__check_allowed, name, value, wrong_dtype=(numbers.Number, Metric, Tensor)
         )
 
         if self._trainer is None:
@@ -492,7 +492,7 @@ def log(
 
     def log_dict(
         self,
-        dictionary: Mapping[str, _METRIC_COLLECTION],
+        dictionary: Mapping[str, _METRIC],
         prog_bar: bool = False,
         logger: Optional[bool] = None,
         on_step: Optional[bool] = None,
@@ -514,8 +514,7 @@ def log_dict(
 
         Args:
             dictionary: key value pairs.
-                The values can be a ``float``, ``Tensor``, ``Metric``, a dictionary of the former
-                or a ``MetricCollection``.
+                The values can be a ``float``, ``Tensor``, ``Metric``, or ``MetricCollection``.
             prog_bar: if ``True`` logs to the progress base.
             logger: if ``True`` logs to the logger.
             on_step: if ``True`` logs at this step.

src/pytorch_lightning/trainer/connectors/logger_connector/result.py

@@ -16,7 +16,7 @@
 from typing import Any, Callable, cast, Dict, Generator, List, Optional, Tuple, Union
 
 import torch
-from lightning_utilities.core.apply_func import apply_to_collection, apply_to_collections
+from lightning_utilities.core.apply_func import apply_to_collection
 from torch import Tensor
 from torchmetrics import Metric
 from typing_extensions import TypedDict
@@ -317,7 +317,6 @@ def __getstate__(self, drop_value: bool = False) -> dict:
             skip.append("value")
         d = {k: v for k, v in self.__dict__.items() if k not in skip}
         d["meta"] = d["meta"].__getstate__()
-        d["_class"] = self.__class__.__name__
         d["_is_synced"] = False  # don't consider the state as synced on reload
         return d
 
@@ -338,48 +337,9 @@ def to(self, *args: Any, **kwargs: Any) -> "_ResultMetric":
         return self
 
 
-class _ResultMetricCollection(dict):
-    """Dict wrapper for easy access to metadata.
-
-    All of the leaf items should be instances of
-    :class:`~pytorch_lightning.trainer.connectors.logger_connector.result._ResultMetric`
-    with the same metadata.
-    """
-
-    @property
-    def meta(self) -> _Metadata:
-        return next(iter(self.values())).meta
-
-    @property
-    def has_tensor(self) -> bool:
-        return any(v.is_tensor for v in self.values())
-
-    def __getstate__(self, drop_value: bool = False) -> dict:
-        def getstate(item: _ResultMetric) -> dict:
-            return item.__getstate__(drop_value=drop_value)
-
-        items = apply_to_collection(dict(self), _ResultMetric, getstate)
-        return {"items": items, "meta": self.meta.__getstate__(), "_class": self.__class__.__name__}
-
-    def __setstate__(self, state: dict, sync_fn: Optional[Callable] = None) -> None:
-        # can't use `apply_to_collection` as it does not recurse items of the same type
-        items = {k: _ResultMetric._reconstruct(v, sync_fn=sync_fn) for k, v in state["items"].items()}
-        self.update(items)
-
-    @classmethod
-    def _reconstruct(cls, state: dict, sync_fn: Optional[Callable] = None) -> "_ResultMetricCollection":
-        rmc = cls()
-        rmc.__setstate__(state, sync_fn=sync_fn)
-        return rmc
-
-
-_METRIC_COLLECTION = Union[_IN_METRIC, _ResultMetricCollection]
-
-
 class _ResultCollection(dict):
-    """
-    Collection (dictionary) of :class:`~pytorch_lightning.trainer.connectors.logger_connector.result._ResultMetric` or
-    :class:`~pytorch_lightning.trainer.connectors.logger_connector.result._ResultMetricCollection`
+    """Collection (dictionary) of
+    :class:`~pytorch_lightning.trainer.connectors.logger_connector.result._ResultMetric`
 
     Example:
 
@@ -404,18 +364,9 @@ def __init__(self, training: bool, device: Optional[Union[str, torch.device]] =
 
     @property
     def result_metrics(self) -> List[_ResultMetric]:
-        o = []
-
-        def append_fn(v: _ResultMetric) -> None:
-            nonlocal o
-            o.append(v)
-
-        apply_to_collection(list(self.values()), _ResultMetric, append_fn)
-        return o
+        return list(self.values())
 
-    def _extract_batch_size(
-        self, value: Union[_ResultMetric, _ResultMetricCollection], batch_size: Optional[int], meta: _Metadata
-    ) -> int:
+    def _extract_batch_size(self, value: _ResultMetric, batch_size: Optional[int], meta: _Metadata) -> int:
         # check if we have extracted the batch size already
         if batch_size is None:
             batch_size = self.batch_size
@@ -424,8 +375,7 @@ def _extract_batch_size(
             return batch_size
 
         batch_size = 1
-        is_tensor = value.is_tensor if isinstance(value, _ResultMetric) else value.has_tensor
-        if self.batch is not None and is_tensor and meta.on_epoch and meta.is_mean_reduction:
+        if self.batch is not None and value.is_tensor and meta.on_epoch and meta.is_mean_reduction:
             batch_size = extract_batch_size(self.batch)
             self.batch_size = batch_size
 
@@ -435,7 +385,7 @@ def log(
         self,
         fx: str,
         name: str,
-        value: _METRIC_COLLECTION,
+        value: _IN_METRIC,
         prog_bar: bool = False,
         logger: bool = True,
         on_step: bool = False,
@@ -494,28 +444,19 @@ def log(
         batch_size = self._extract_batch_size(self[key], batch_size, meta)
         self.update_metrics(key, value, batch_size)
 
-    def register_key(self, key: str, meta: _Metadata, value: _METRIC_COLLECTION) -> None:
+    def register_key(self, key: str, meta: _Metadata, value: _IN_METRIC) -> None:
         """Create one _ResultMetric object per value.
 
         Value can be provided as a nested collection
         """
+        metric = _ResultMetric(meta, isinstance(value, Tensor)).to(self.device)
+        self[key] = metric
 
-        def fn(v: _IN_METRIC) -> _ResultMetric:
-            metric = _ResultMetric(meta, isinstance(v, Tensor))
-            return metric.to(self.device)
-
-        value = apply_to_collection(value, (Tensor, Metric), fn)
-        if isinstance(value, dict):
-            value = _ResultMetricCollection(value)
-        self[key] = value
-
-    def update_metrics(self, key: str, value: _METRIC_COLLECTION, batch_size: int) -> None:
-        def fn(result_metric: _ResultMetric, v: Tensor) -> None:
-            # performance: avoid calling `__call__` to avoid the checks in `torch.nn.Module._call_impl`
-            result_metric.forward(v.to(self.device), batch_size)
-            result_metric.has_reset = False
-
-        apply_to_collections(self[key], value, _ResultMetric, fn)
+    def update_metrics(self, key: str, value: _IN_METRIC, batch_size: int) -> None:
+        result_metric = self[key]
+        # performance: avoid calling `__call__` to avoid the checks in `torch.nn.Module._call_impl`
+        result_metric.forward(value.to(self.device), batch_size)
+        result_metric.has_reset = False
 
     @staticmethod
     def _get_cache(result_metric: _ResultMetric, on_step: bool) -> Optional[Tensor]:
@@ -557,11 +498,7 @@ def _get_cache(result_metric: _ResultMetric, on_step: bool) -> Optional[Tensor]:
 
     def valid_items(self) -> Generator:
         """This function is used to iterate over current valid metrics."""
-        return (
-            (k, v)
-            for k, v in self.items()
-            if not (isinstance(v, _ResultMetric) and v.has_reset) and self.dataloader_idx == v.meta.dataloader_idx
-        )
+        return ((k, v) for k, v in self.items() if not v.has_reset and self.dataloader_idx == v.meta.dataloader_idx)
 
     def _forked_name(self, result_metric: _ResultMetric, on_step: bool) -> Tuple[str, str]:
         name = result_metric.meta.name
@@ -578,23 +515,9 @@ def metrics(self, on_step: bool) -> _METRICS:
         metrics = _METRICS(callback={}, log={}, pbar={})
 
         for _, result_metric in self.valid_items():
-
-            # extract forward_cache or computed from the _ResultMetric. ignore when the output is None
-            value = apply_to_collection(result_metric, _ResultMetric, self._get_cache, on_step, include_none=False)
-
-            # convert metric collection to dict container.
-            if isinstance(value, _ResultMetricCollection):
-                value = dict(value.items())
-
-            # check if the collection is empty
-            has_tensor = False
-
-            def any_tensor(_: Any) -> None:
-                nonlocal has_tensor
-                has_tensor = True
-
-            apply_to_collection(value, Tensor, any_tensor)
-            if not has_tensor:
+            # extract forward_cache or computed from the _ResultMetric
+            value = self._get_cache(result_metric, on_step)
+            if not isinstance(value, Tensor):
                 continue
 
             name, forked_name = self._forked_name(result_metric, on_step)
@@ -623,15 +546,12 @@ def reset(self, metrics: Optional[bool] = None, fx: Optional[str] = None) -> Non
                 if ``None``, both are.
             fx: Function to reset
         """
-
-        def fn(item: _ResultMetric) -> None:
+        for item in self.values():
             requested_type = metrics is None or metrics ^ item.is_tensor
             same_fx = fx is None or fx == item.meta.fx
             if requested_type and same_fx:
                 item.reset()
 
-        apply_to_collection(self, _ResultMetric, fn)
-
     def to(self, *args: Any, **kwargs: Any) -> "_ResultCollection":
         """Move all data to the given device."""
         self.update(apply_to_collection(dict(self), (Tensor, Metric), move_data_to_device, *args, **kwargs))
@@ -664,7 +584,6 @@ def __repr__(self) -> str:
 
     def __getstate__(self, drop_value: bool = True) -> dict:
         d = self.__dict__.copy()
-        # all the items should be either `_ResultMetric`s or `_ResultMetricCollection`s
         items = {k: v.__getstate__(drop_value=drop_value) for k, v in self.items()}
         return {**d, "items": items}
 
@@ -673,18 +592,11 @@ def __setstate__(
     ) -> None:
         self.__dict__.update({k: v for k, v in state.items() if k != "items"})
 
-        def setstate(k: str, item: dict) -> Union[_ResultMetric, _ResultMetricCollection]:
+        def setstate(k: str, item: dict) -> _ResultMetric:
             if not isinstance(item, dict):
                 raise ValueError(f"Unexpected value: {item}")
-            cls = item["_class"]
-            if cls == _ResultMetric.__name__:
-                cls = _ResultMetric
-            elif cls == _ResultMetricCollection.__name__:
-                cls = _ResultMetricCollection
-            else:
-                raise ValueError(f"Unexpected class name: {cls}")
             _sync_fn = sync_fn or (self[k].meta.sync.fn if k in self else None)
-            return cls._reconstruct(item, sync_fn=_sync_fn)
+            return _ResultMetric._reconstruct(item, sync_fn=_sync_fn)
 
         items = {k: setstate(k, v) for k, v in state["items"].items()}
         self.update(items)

src/pytorch_lightning/utilities/types.py

@@ -19,7 +19,7 @@
 from argparse import _ArgumentGroup, ArgumentParser
 from contextlib import contextmanager
 from dataclasses import dataclass
-from typing import Any, Dict, Generator, List, Mapping, Optional, Sequence, Type, Union
+from typing import Any, Dict, Generator, List, Optional, Sequence, Type, Union
 
 import torch
 from torch import Tensor
@@ -31,7 +31,6 @@
 
 _NUMBER = Union[int, float]
 _METRIC = Union[Metric, Tensor, _NUMBER]
-_METRIC_COLLECTION = Union[_METRIC, Mapping[str, _METRIC]]
 STEP_OUTPUT = Union[Tensor, Dict[str, Any]]
 EPOCH_OUTPUT = List[STEP_OUTPUT]
 _EVALUATE_OUTPUT = List[Dict[str, float]]  # 1 dict per DataLoader

tests/tests_pytorch/callbacks/progress/test_tqdm_progress_bar.py

@@ -407,8 +407,8 @@ def test_tensor_to_float_conversion(tmpdir):
     class TestModel(BoringModel):
         def training_step(self, batch, batch_idx):
             self.log("a", torch.tensor(0.123), prog_bar=True, on_epoch=False)
-            self.log("b", {"b1": torch.tensor([1])}, prog_bar=True, on_epoch=False)
-            self.log("c", {"c1": 2}, prog_bar=True, on_epoch=False)
+            self.log("b", torch.tensor([1]), prog_bar=True, on_epoch=False)
+            self.log("c", 2, prog_bar=True, on_epoch=False)
             return super().training_step(batch, batch_idx)
 
     trainer = Trainer(
@@ -417,11 +417,11 @@ def training_step(self, batch, batch_idx):
     trainer.fit(TestModel())
 
     torch.testing.assert_close(trainer.progress_bar_metrics["a"], 0.123)
-    assert trainer.progress_bar_metrics["b"] == {"b1": 1.0}
-    assert trainer.progress_bar_metrics["c"] == {"c1": 2.0}
+    assert trainer.progress_bar_metrics["b"] == 1.0
+    assert trainer.progress_bar_metrics["c"] == 2.0
     pbar = trainer.progress_bar_callback.main_progress_bar
     actual = str(pbar.postfix)
-    assert actual.endswith("a=0.123, b={'b1': 1.0}, c={'c1': 2.0}"), actual
+    assert actual.endswith("a=0.123, b=1.000, c=2.000"), actual
 
 
 @pytest.mark.parametrize(

tests/tests_pytorch/core/test_metric_result_integration.py

@@ -237,11 +237,10 @@ def lightning_log(fx, *args, **kwargs):
             lightning_log("training_step", "c", metric_c, on_step=True, on_epoch=False, metric_attribute="metric_c")
             lightning_log("training_step", "a_1", a, on_step=True, on_epoch=True)
             lightning_log("training_step", "b_1", b, on_step=False, on_epoch=True)
-            lightning_log("training_step", "c_1", {"1": c, "2": c}, on_step=True, on_epoch=False)
+            lightning_log("training_step", "c_1", c, on_step=True, on_epoch=False)
 
             batch_log = result.metrics(on_step=True)["log"]
             assert set(batch_log) == {"a_step", "c", "a_1_step", "c_1"}
-            assert set(batch_log["c_1"]) == {"1", "2"}
 
             result_copy = deepcopy(result)
             new_result = _ResultCollection(True, torch.device("cpu"))
@@ -250,7 +249,7 @@ def lightning_log(fx, *args, **kwargs):
             assert "fn" not in state_dict["items"]["training_step.a"]["meta"]["_sync"]
 
             assert not new_result.result_metrics
-            assert len(result.result_metrics) == 7 + epoch > 0
+            assert len(result.result_metrics) == 6 + epoch > 0
 
             new_result.load_state_dict(
                 state_dict, metrics={"metric": metric, "metric_b": metric_b, "metric_c": metric_c}