Add typing for ResultCollection [3/3] (#9271)

Author: carmocca

pyproject.toml

@@ -64,8 +64,7 @@ module = [
     "pytorch_lightning.callbacks.pruning",
     "pytorch_lightning.loops.closure",
     "pytorch_lightning.trainer.evaluation_loop",
-    "pytorch_lightning.trainer.connectors.logger_connector.fx_validator",
-    "pytorch_lightning.trainer.connectors.logger_connector.logger_connector",
+    "pytorch_lightning.trainer.connectors.logger_connector.*",
     "pytorch_lightning.trainer.progress",
     "pytorch_lightning.tuner.auto_gpu_select",
     "pytorch_lightning.utilities.apply_func",

pytorch_lightning/trainer/connectors/logger_connector/result.py

@@ -14,7 +14,7 @@
 from collections.abc import Generator
 from dataclasses import asdict, dataclass, replace
 from functools import partial, wraps
-from typing import Any, Callable, Dict, List, Mapping, Optional, Tuple, Union
+from typing import Any, Callable, Dict, List, Optional, Tuple, Union
 
 import torch
 from torchmetrics import Metric
@@ -24,16 +24,13 @@
 from pytorch_lightning.utilities import rank_zero_warn
 from pytorch_lightning.utilities.apply_func import apply_to_collection, apply_to_collections, move_data_to_device
 from pytorch_lightning.utilities.data import extract_batch_size
-from pytorch_lightning.utilities.enums import LightningEnum
 from pytorch_lightning.utilities.exceptions import MisconfigurationException
 from pytorch_lightning.utilities.memory import recursive_detach
 from pytorch_lightning.utilities.metrics import metrics_to_scalars
 from pytorch_lightning.utilities.warnings import WarningCache
 
-# re-define the ones from pytorch_lightning.utilities.types without the `Number` type
 # TODO(@tchaton): Typing-pickle issue on python<3.7 (https://github.com/cloudpipe/cloudpickle/pull/318)
-_METRIC = Any  # Union[Metric, torch.Tensor]
-_METRIC_COLLECTION = Union[_METRIC, Mapping[str, _METRIC]]
+_IN_METRIC = Any  # Union[Metric, torch.Tensor]  # Do not include scalars as they were converted to tensors
 _OUT_METRIC = Union[torch.Tensor, Dict[str, torch.Tensor]]
 _PBAR_METRIC = Union[float, Dict[str, float]]
 _OUT_DICT = Dict[str, _OUT_METRIC]
@@ -49,12 +46,6 @@ class _METRICS(TypedDict):
 warning_cache = WarningCache()
 
 
-class MetricSource(LightningEnum):
-    CALLBACK = "callback"
-    PBAR = "pbar"
-    LOG = "log"
-
-
 @dataclass
 class _Sync:
     fn: Optional[Callable] = None
@@ -80,14 +71,15 @@ def _generate_sync_fn(self) -> None:
         """Used to compute the syncing function and cache it."""
         fn = self.no_op if self.fn is None or not self.should or self.rank_zero_only else self.fn
         # save the function as `_fn` as the meta are being re-created and the object references need to match.
-        self._fn = partial(fn, reduce_op=self.op, group=self.group)
+        # ignore typing, bad support for `partial`: mypy/issues/1484
+        self._fn: Callable = partial(fn, reduce_op=self.op, group=self.group)  # type: ignore [arg-type]
 
     @property
     def __call__(self) -> Any:
         return self._fn
 
     @staticmethod
-    def no_op(value: Any, *_, **__) -> Any:
+    def no_op(value: Any, *_: Any, **__: Any) -> Any:
         return value
 
 
@@ -125,7 +117,8 @@ def _parse_reduce_fx(self) -> None:
             raise MisconfigurationException(error)
 
     @property
-    def sync(self) -> Optional[_Sync]:
+    def sync(self) -> _Sync:
+        assert self._sync is not None
         return self._sync
 
     @sync.setter
@@ -196,7 +189,7 @@ def __init__(self, metadata: _Metadata, is_tensor: bool) -> None:
             if self.meta.is_mean_reduction:
                 self.add_state("cumulated_batch_size", torch.tensor(0, dtype=torch.float), dist_reduce_fx=torch.sum)
 
-    def update(self, value: _METRIC, batch_size: torch.Tensor) -> None:
+    def update(self, value: _IN_METRIC, batch_size: torch.Tensor) -> None:
         if self.is_tensor:
             value = value.float()
             # performance: no need to accumulate on values only logged on_step
@@ -232,7 +225,7 @@ def reset(self) -> None:
             self.value.reset()
         self.has_reset = True
 
-    def forward(self, value: _METRIC, batch_size: torch.Tensor) -> None:
+    def forward(self, value: _IN_METRIC, batch_size: torch.Tensor) -> None:
         if self.meta.enable_graph:
             with torch.no_grad():
                 self.update(value, batch_size)
@@ -243,7 +236,7 @@ def forward(self, value: _METRIC, batch_size: torch.Tensor) -> None:
     def _wrap_compute(self, compute: Any) -> Any:
         # Override to avoid syncing - we handle it ourselves.
         @wraps(compute)
-        def wrapped_func(*args, **kwargs):
+        def wrapped_func(*args: Any, **kwargs: Any) -> Optional[Any]:
             if not self._update_called:
                 rank_zero_warn(
                     f"The ``compute`` method of metric {self.__class__.__name__}"
@@ -253,8 +246,8 @@ def wrapped_func(*args, **kwargs):
                 )
 
             # return cached value
-            if self._computed is not None:
-                return self._computed
+            if self._computed is not None:  # type: ignore
+                return self._computed  # type: ignore
             self._computed = compute(*args, **kwargs)
             return self._computed
 
@@ -293,7 +286,7 @@ def _reconstruct(cls, state: dict, sync_fn: Optional[Callable] = None) -> "Resul
         result_metric.__setstate__(state, sync_fn=sync_fn)
         return result_metric
 
-    def to(self, *args: Any, **kwargs: Any) -> "DeviceDtypeModuleMixin":
+    def to(self, *args: Any, **kwargs: Any) -> "ResultMetric":
         self.__dict__.update(
             apply_to_collection(self.__dict__, (torch.Tensor, Metric), move_data_to_device, *args, **kwargs)
         )
@@ -309,7 +302,7 @@ class ResultMetricCollection(dict):
     with the same metadata.
     """
 
-    def __init__(self, *args) -> None:
+    def __init__(self, *args: Any) -> None:
         super().__init__(*args)
 
     @property
@@ -320,20 +313,12 @@ 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, ResultMetricCollection), getstate)
+        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:
-        def setstate(item: dict) -> Union[Dict[str, ResultMetric], ResultMetric, Any]:
-            # recurse through dictionaries to set the state. can't use `apply_to_collection`
-            # as it does not recurse items of the same type.
-            if not isinstance(item, dict):
-                return item
-            if item.get("_class") == ResultMetric.__name__:
-                return ResultMetric._reconstruct(item, sync_fn=sync_fn)
-            return {k: setstate(v) for k, v in item.items()}
-
-        items = setstate(state["items"])
+        # 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
@@ -343,6 +328,9 @@ def _reconstruct(cls, state: dict, sync_fn: Optional[Callable] = None) -> "Resul
         return rmc
 
 
+_METRIC_COLLECTION = Union[_IN_METRIC, ResultMetricCollection]
+
+
 class ResultCollection(dict):
     """
     Collection (dictionary) of :class:`~pytorch_lightning.trainer.connectors.logger_connector.result.ResultMetric` or
@@ -364,7 +352,7 @@ class ResultCollection(dict):
     def __init__(self, training: bool, device: Optional[Union[str, torch.device]] = None) -> None:
         super().__init__()
         self.training = training
-        self._minimize = None
+        self._minimize: Optional[torch.Tensor] = None
         self._batch_size = torch.tensor(1, device=device)
         self.device: Optional[Union[str, torch.device]] = device
 
@@ -413,7 +401,7 @@ def extra(self) -> Dict[str, Any]:
 
     @extra.setter
     def extra(self, extra: Dict[str, Any]) -> None:
-        def check_fn(v):
+        def check_fn(v: torch.Tensor) -> torch.Tensor:
             if v.grad_fn is not None:
                 warning_cache.deprecation(
                     f"One of the returned values {set(extra.keys())} has a `grad_fn`. We will detach it automatically"
@@ -494,7 +482,7 @@ def log(
     def register_key(self, key: str, meta: _Metadata, value: _METRIC_COLLECTION) -> None:
         """Create one ResultMetric object per value. Value can be provided as a nested collection"""
 
-        def fn(v: _METRIC) -> ResultMetric:
+        def fn(v: _IN_METRIC) -> ResultMetric:
             metric = ResultMetric(meta, isinstance(v, torch.Tensor))
             return metric.to(self.device)
 
@@ -504,7 +492,7 @@ def fn(v: _METRIC) -> ResultMetric:
         self[key] = value
 
     def update_metrics(self, key: str, value: _METRIC_COLLECTION) -> None:
-        def fn(result_metric, v):
+        def fn(result_metric: ResultMetric, v: ResultMetric) -> None:
             # performance: avoid calling `__call__` to avoid the checks in `torch.nn.Module._call_impl`
             result_metric.forward(v.to(self.device), self.batch_size)
             result_metric.has_reset = False
@@ -545,7 +533,7 @@ def _forked_name(self, result_metric: ResultMetric, on_step: bool) -> Tuple[str,
         return name, forked_name
 
     def metrics(self, on_step: bool) -> _METRICS:
-        metrics = {k: {} for k in MetricSource}
+        metrics = _METRICS(callback={}, log={}, pbar={})
 
         for _, result_metric in self.valid_items():
 
@@ -559,7 +547,7 @@ def metrics(self, on_step: bool) -> _METRICS:
             # check if the collection is empty
             has_tensor = False
 
-            def any_tensor(_):
+            def any_tensor(_: Any) -> None:
                 nonlocal has_tensor
                 has_tensor = True
 
@@ -571,16 +559,16 @@ def any_tensor(_):
 
             # populate logging metrics
             if result_metric.meta.logger:
-                metrics[MetricSource.LOG][forked_name] = value
+                metrics["log"][forked_name] = value
 
             # populate callback metrics. callback metrics don't take `_step` forked metrics
             if self.training or result_metric.meta.on_epoch and not on_step:
-                metrics[MetricSource.CALLBACK][name] = value
-                metrics[MetricSource.CALLBACK][forked_name] = value
+                metrics["callback"][name] = value
+                metrics["callback"][forked_name] = value
 
             # populate progress_bar metrics. convert tensors to numbers
             if result_metric.meta.prog_bar:
-                metrics[MetricSource.PBAR][forked_name] = metrics_to_scalars(value)
+                metrics["pbar"][forked_name] = metrics_to_scalars(value)
 
         return metrics
 
@@ -609,7 +597,7 @@ def extract_batch_size(self, batch: Any) -> None:
         except RecursionError:
             self.batch_size = 1
 
-    def to(self, *args, **kwargs) -> "ResultCollection":
+    def to(self, *args: Any, **kwargs: Any) -> "ResultCollection":
         """Move all data to the given device."""
         self.update(apply_to_collection(dict(self), (torch.Tensor, Metric), move_data_to_device, *args, **kwargs))
 
@@ -641,7 +629,7 @@ def __str__(self) -> str:
         self_str = str({k: v for k, v in self.items() if v})
         return f"{self.__class__.__name__}({minimize}{self_str})"
 
-    def __repr__(self):
+    def __repr__(self) -> str:
         # sample output: `{True, cpu, minimize=tensor(1.23 grad_fn=<SumBackward0>), {'_extra': {}}}`
         minimize = f"minimize={repr(self.minimize)}, " if self.minimize is not None else ""
         return f"{{{self.training}, {repr(self.device)}, " + minimize + f"{super().__repr__()}}}"

pytorch_lightning/utilities/types.py

@@ -16,7 +16,6 @@
  - Do not include any `_TYPE` suffix
  - Types used in public hooks (as those in the `LightningModule` and `Callback`) should be public (no trailing `_`)
 """
-from numbers import Number
 from pathlib import Path
 from typing import Any, Dict, Iterator, List, Mapping, Sequence, Type, Union
 
@@ -25,7 +24,8 @@
 from torch.utils.data import DataLoader
 from torchmetrics import Metric
 
-_METRIC = Union[Metric, torch.Tensor, Number]
+_NUMBER = Union[int, float]
+_METRIC = Union[Metric, torch.Tensor, _NUMBER]
 _METRIC_COLLECTION = Union[_METRIC, Mapping[str, _METRIC]]
 STEP_OUTPUT = Union[torch.Tensor, Dict[str, Any]]
 EPOCH_OUTPUT = List[STEP_OUTPUT]

tests/core/test_metric_result_integration.py

@@ -27,7 +27,7 @@
 import tests.helpers.utils as tutils
 from pytorch_lightning import Trainer
 from pytorch_lightning.callbacks import ModelCheckpoint
-from pytorch_lightning.trainer.connectors.logger_connector.result import MetricSource, ResultCollection
+from pytorch_lightning.trainer.connectors.logger_connector.result import ResultCollection
 from pytorch_lightning.utilities.imports import _fault_tolerant_training, _TORCH_GREATER_EQUAL_1_7
 from tests.helpers import BoringModel
 from tests.helpers.runif import RunIf
@@ -81,10 +81,10 @@ def _ddp_test_fn(rank, worldsize):
             result.log("h", "b", metric_b, on_step=False, on_epoch=True)
             result.log("h", "c", metric_c, on_step=True, on_epoch=False)
 
-            batch_log = result.metrics(True)[MetricSource.LOG]
+            batch_log = result.metrics(True)["log"]
             assert batch_log == {"a_step": i, "c": i}
 
-        epoch_log = result.metrics(False)[MetricSource.LOG]
+        epoch_log = result.metrics(False)["log"]
         result.reset()
 
         # assert metric state reset to default values
@@ -124,10 +124,10 @@ def test_result_metric_integration():
             result.log("h", "b", metric_b, on_step=False, on_epoch=True)
             result.log("h", "c", metric_c, on_step=True, on_epoch=False)
 
-            batch_log = result.metrics(True)[MetricSource.LOG]
+            batch_log = result.metrics(True)["log"]
             assert batch_log == {"a_step": i, "c": i}
 
-        epoch_log = result.metrics(False)[MetricSource.LOG]
+        epoch_log = result.metrics(False)["log"]
         result.reset()
 
         # assert metric state reset to default values
@@ -248,7 +248,7 @@ def lightning_log(fx, *args, **kwargs):
             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)
 
-            batch_log = result.metrics(on_step=True)[MetricSource.LOG]
+            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"}
 
@@ -269,12 +269,12 @@ def lightning_log(fx, *args, **kwargs):
             # the sync fn has been kept
             assert result_copy["training_step.a"].meta.sync.fn == new_result["training_step.a"].meta.sync.fn
 
-        epoch_log = result.metrics(on_step=False)[MetricSource.LOG]
-        epoch_log_copy = result_copy.metrics(on_step=False)[MetricSource.LOG]
+        epoch_log = result.metrics(on_step=False)["log"]
+        epoch_log_copy = result_copy.metrics(on_step=False)["log"]
         assert epoch_log == epoch_log_copy
 
         lightning_log("train_epoch_end", "a", metric_a, on_step=False, on_epoch=True)
-        epoch_log = result.metrics(on_step=False)[MetricSource.LOG]
+        epoch_log = result.metrics(on_step=False)["log"]
         assert epoch_log == {
             "a_1_epoch": 1,
             "a_epoch": cumulative_sum,
@@ -451,9 +451,9 @@ def on_epoch_end(self) -> None:
                 total = sum(range(5)) * num_processes
                 metrics = self.results.metrics(on_step=False)
                 assert self.results["training_step.tracking"].value == total
-                assert metrics[MetricSource.CALLBACK]["tracking"] == self.dummy_metric.compute() == 2
+                assert metrics["callback"]["tracking"] == self.dummy_metric.compute() == 2
                 assert self.results["training_step.tracking_2"].value == total
-                assert metrics[MetricSource.CALLBACK]["tracking_2"] == self.dummy_metric.compute() == 2
+                assert metrics["callback"]["tracking_2"] == self.dummy_metric.compute() == 2
                 self.has_validated_sum = True
 
     model = ExtendedBoringModel()