diff --git a/ignite/metrics/clustering/_base.py b/ignite/metrics/clustering/_base.py index 9650e9ed62ed..ad51e33587c5 100644 --- a/ignite/metrics/clustering/_base.py +++ b/ignite/metrics/clustering/_base.py @@ -1,3 +1,5 @@ +from typing import cast + from torch import Tensor from ignite.exceptions import NotComputableError from ignite.metrics.epoch_metric import EpochMetric @@ -36,4 +38,4 @@ def compute(self) -> float: f"{self.__class__.__name__} must have at least one example before it can be computed." ) - return super().compute() + return cast(float, super().compute()) diff --git a/ignite/metrics/clustering/calinski_harabasz_score.py b/ignite/metrics/clustering/calinski_harabasz_score.py index 8db162c337cc..6293442021f8 100644 --- a/ignite/metrics/clustering/calinski_harabasz_score.py +++ b/ignite/metrics/clustering/calinski_harabasz_score.py @@ -20,7 +20,7 @@ def _calinski_harabasz_score(features: Tensor, labels: Tensor) -> float: np_features = features.cpu().numpy() np_labels = labels.cpu().numpy() score = calinski_harabasz_score(np_features, np_labels) - return score + return float(score) class CalinskiHarabaszScore(_ClusteringMetricBase): diff --git a/ignite/metrics/clustering/davies_bouldin_score.py b/ignite/metrics/clustering/davies_bouldin_score.py index af448eaf60f5..1922e985dae8 100644 --- a/ignite/metrics/clustering/davies_bouldin_score.py +++ b/ignite/metrics/clustering/davies_bouldin_score.py @@ -20,7 +20,7 @@ def _davies_bouldin_score(features: Tensor, labels: Tensor) -> float: np_features = features.cpu().numpy() np_labels = labels.cpu().numpy() score = davies_bouldin_score(np_features, np_labels) - return score + return float(score) class DaviesBouldinScore(_ClusteringMetricBase): diff --git a/ignite/metrics/clustering/silhouette_score.py b/ignite/metrics/clustering/silhouette_score.py index 2acb9c25e38d..3c290621196e 100644 --- a/ignite/metrics/clustering/silhouette_score.py +++ b/ignite/metrics/clustering/silhouette_score.py @@ -127,4 +127,4 @@ def _silhouette_score(self, features: Tensor, labels: Tensor) -> float: np_labels = labels.detach().cpu().numpy() score = silhouette_score(np_features, np_labels, **self._silhouette_kwargs) - return score + return float(score) diff --git a/ignite/metrics/epoch_metric.py b/ignite/metrics/epoch_metric.py index 672f843d1697..fd7debb4612d 100644 --- a/ignite/metrics/epoch_metric.py +++ b/ignite/metrics/epoch_metric.py @@ -1,6 +1,6 @@ import warnings -from collections.abc import Callable -from typing import cast +from collections.abc import Callable, Mapping, Sequence +from typing import Union, cast import torch @@ -8,8 +8,22 @@ from ignite.exceptions import NotComputableError from ignite.metrics.metric import Metric, reinit__is_reduced +# Supported return types for ``EpochMetric``'s ``compute_fn``. +EpochMetricOutput = Union[int, float, torch.Tensor, Sequence, Mapping] + __all__ = ["EpochMetric"] +# Type tags used by ``EpochMetric._broadcast_result`` to let every rank agree on how to +# decode the next value(s) coming out of the collective calls. ``_TAG_UNSUPPORTED`` is also +# reused to reject non-str mapping keys, since both cases mean "every rank should raise". +_TAG_INT = 0 +_TAG_FLOAT = 1 +_TAG_TENSOR = 2 +_TAG_TUPLE = 3 +_TAG_LIST = 4 +_TAG_MAPPING = 5 +_TAG_UNSUPPORTED = -1 + class EpochMetric(Metric): """Class for metrics that should be computed on the entire output history of a model. @@ -30,7 +44,13 @@ class EpochMetric(Metric): Args: compute_fn: a callable which receives two tensors as the `predictions` and `targets` - and returns a scalar. Input tensors will be on specified ``device`` (see arg below). + and returns the computed metric. Supported return types are: ``int``, ``float``, + ``torch.Tensor``, a ``Sequence`` (tuple/list) of these, or a ``Mapping`` (dict) with + string keys and these values, including arbitrarily nested combinations of these. + An unsupported return type raises a ``TypeError``. These types are also supported + in distributed configuration (``world_size > 1``): the result is broadcast from + rank 0 to all other ranks. Input tensors will be on specified ``device`` + (see arg below). output_transform: a callable that is used to transform the :class:`~ignite.engine.engine.Engine`'s ``process_function``'s output into the form expected by the metric. This can be useful if, for example, you have a multi-output model and @@ -93,7 +113,7 @@ def __init__( def reset(self) -> None: self._predictions: list[torch.Tensor] = [] self._targets: list[torch.Tensor] = [] - self._result: float | None = None + self._result: EpochMetricOutput | None = None def _check_shape(self, output: tuple[torch.Tensor, torch.Tensor]) -> None: y_pred, y = output @@ -142,7 +162,114 @@ def update(self, output: tuple[torch.Tensor, torch.Tensor]) -> None: except Exception as e: warnings.warn(f"Probably, there can be a problem with `compute_fn`:\n {e}.", EpochMetricWarning) - def compute(self) -> float: + def _check_output_type(self, result: EpochMetricOutput) -> None: + # Recursively validate that compute_fn's output is a supported type. ``str``/``bytes`` + # are rejected explicitly since ``str`` is itself a ``Sequence``. + if isinstance(result, (int, float, torch.Tensor)): + return + if isinstance(result, Mapping): + for key, value in result.items(): + if not isinstance(key, str): + raise TypeError(f"compute_fn output mapping keys should be str, but given {type(key)}.") + self._check_output_type(value) + return + if isinstance(result, Sequence) and not isinstance(result, (str, bytes)): + for value in result: + self._check_output_type(value) + return + raise TypeError( + f"compute_fn output type {type(result)} is not supported. Supported types are: " + "int, float, torch.Tensor, a Sequence of these, or a Mapping with str keys and these values." + ) + + def _broadcast_result(self, result: EpochMetricOutput, src: int = 0) -> EpochMetricOutput: + """Recursively broadcast compute_fn output from src rank to all ranks. + + Each step only broadcasts types that ``idist.broadcast`` natively supports + (int, float, torch.Tensor, str), so containers are transmitted by first + synchronising their structure (type tag, length, dict keys) and then + broadcasting each leaf individually. + + Every rank must take the same path through the collective calls below, so any + rejection (unsupported type, non-str mapping key) is decided from a value that has + already been broadcast to all ranks, never from a check that only src has run. This + way every rank raises together instead of some ranks hanging on a broadcast that src + never issues. + """ + rank = idist.get_rank() + + # Step 1: broadcast type tag so all ranks know what to expect + if rank == src: + if isinstance(result, int): + tag = _TAG_INT + elif isinstance(result, float): + tag = _TAG_FLOAT + elif isinstance(result, torch.Tensor): + tag = _TAG_TENSOR + elif isinstance(result, tuple): + tag = _TAG_TUPLE + elif isinstance(result, list): + tag = _TAG_LIST + elif isinstance(result, Mapping): + tag = _TAG_MAPPING + else: + tag = _TAG_UNSUPPORTED + else: + tag = _TAG_INT + tag = cast(int, idist.broadcast(tag, src=src)) + + if tag == _TAG_UNSUPPORTED: + raise TypeError( + "compute_fn output type is not supported. Supported types are: " + "int, float, torch.Tensor, a Sequence of these, or a Mapping with str keys and these values." + ) + + # Step 2: broadcast content based on type. ``result`` is only meaningfully typed on + # ``src`` (the tag protocol above guarantees every rank agrees on which branch runs), + # so it is cast to the type the tag promises; non-src ranks only need a same-typed + # placeholder since their value is discarded by the collective call. + if tag == _TAG_INT: + int_value = cast(int, result) if rank == src else 0 + return cast(int, idist.broadcast(int_value, src=src)) + if tag == _TAG_FLOAT: + float_value = cast(float, result) if rank == src else 0.0 + return cast(float, idist.broadcast(float_value, src=src)) + if tag == _TAG_TENSOR: + tensor_value = cast(torch.Tensor, result) if rank == src else None + return cast(torch.Tensor, idist.broadcast(tensor_value, src=src, safe_mode=True)) + + if tag in (_TAG_TUPLE, _TAG_LIST): + seq_value = cast(Sequence, result) if rank == src else [] + length = cast(int, idist.broadcast(len(seq_value) if rank == src else 0, src=src)) + elements = [] + for i in range(length): + elem = cast(EpochMetricOutput, seq_value[i]) if rank == src else 0 + elements.append(self._broadcast_result(elem, src=src)) + return tuple(elements) if tag == _TAG_TUPLE else elements + + # tag == _TAG_MAPPING + mapping_value = cast(Mapping, result) if rank == src else {} + src_keys = list(mapping_value.keys()) if rank == src else [] + n_keys = cast(int, idist.broadcast(len(src_keys) if rank == src else 0, src=src)) + keys = [] + for i in range(n_keys): + raw_key = src_keys[i] if rank == src else None + # Broadcast whether this key is a valid (str) key before broadcasting the key + # itself: src and non-src ranks must agree on the wire type (str) they are about + # to exchange, so this can't be decided from a src-only isinstance check. + key_valid = cast(int, idist.broadcast(0 if (rank != src or isinstance(raw_key, str)) else -1, src=src)) + if key_valid == _TAG_UNSUPPORTED: + detail = f" but given {type(raw_key)}" if rank == src else "" + raise TypeError(f"compute_fn output mapping keys should be str{detail}.") + key = cast(str, idist.broadcast(raw_key if rank == src else "", src=src)) + keys.append(key) + values = [] + for i in range(n_keys): + val = cast(EpochMetricOutput, mapping_value[keys[i]]) if rank == src else 0 + values.append(self._broadcast_result(val, src=src)) + return dict(zip(keys, values)) + + def compute(self) -> EpochMetricOutput: if len(self._predictions) < 1 or len(self._targets) < 1: raise NotComputableError(f"{type(self).__name__} must have at least one example before it can be computed.") @@ -156,14 +283,22 @@ def compute(self) -> float: _prediction_tensor = cast(torch.Tensor, idist.all_gather(_prediction_tensor)) _target_tensor = cast(torch.Tensor, idist.all_gather(_target_tensor)) - self._result = 0.0 + result: EpochMetricOutput = 0.0 if idist.get_rank() == 0: # Run compute_fn on zero rank only - self._result = self.compute_fn(_prediction_tensor, _target_tensor) + result = self.compute_fn(_prediction_tensor, _target_tensor) if ws > 1: - # broadcast result to all processes - self._result = cast(float, idist.broadcast(self._result, src=0)) + # Type/key validation happens inside `_broadcast_result` itself (via the tag + # protocol), so every rank reaches the same TypeError together. Do not + # pre-validate on rank 0 alone here: that would let rank 0 raise and return + # before issuing the first broadcast, leaving other ranks waiting on a + # collective call rank 0 never makes. + result = self._broadcast_result(result, src=0) + else: + self._check_output_type(result) + + self._result = result return self._result diff --git a/ignite/metrics/regression/kendall_correlation.py b/ignite/metrics/regression/kendall_correlation.py index d47353162a7e..5ca9ffcffd6c 100644 --- a/ignite/metrics/regression/kendall_correlation.py +++ b/ignite/metrics/regression/kendall_correlation.py @@ -1,5 +1,5 @@ from collections.abc import Callable -from typing import Any +from typing import Any, cast import torch @@ -20,7 +20,7 @@ def _tau(predictions: Tensor, targets: Tensor) -> float: np_preds = predictions.flatten().cpu().numpy() np_targets = targets.flatten().cpu().numpy() r = kendalltau(np_preds, np_targets, variant=variant).statistic - return r + return float(r) return _tau @@ -121,4 +121,4 @@ def compute(self) -> float: if len(self._predictions) < 1 or len(self._targets) < 1: raise NotComputableError("KendallRankCorrelation must have at least one example before it can be computed.") - return super().compute() + return cast(float, super().compute()) diff --git a/ignite/metrics/regression/spearman_correlation.py b/ignite/metrics/regression/spearman_correlation.py index 755698b6f867..292006c38931 100644 --- a/ignite/metrics/regression/spearman_correlation.py +++ b/ignite/metrics/regression/spearman_correlation.py @@ -1,4 +1,4 @@ -from typing import Any +from typing import Any, cast from collections.abc import Callable import torch @@ -16,7 +16,7 @@ def _spearman_r(predictions: Tensor, targets: Tensor) -> float: np_preds = predictions.flatten().cpu().numpy() np_targets = targets.flatten().cpu().numpy() r = spearmanr(np_preds, np_targets).statistic - return r + return float(r) class SpearmanRankCorrelation(EpochMetric): @@ -110,4 +110,4 @@ def compute(self) -> float: "SpearmanRankCorrelation must have at least one example before it can be computed." ) - return super().compute() + return cast(float, super().compute()) diff --git a/tests/ignite/metrics/test_epoch_metric.py b/tests/ignite/metrics/test_epoch_metric.py index 5bbb2e2307cc..1911eb12eeb3 100644 --- a/tests/ignite/metrics/test_epoch_metric.py +++ b/tests/ignite/metrics/test_epoch_metric.py @@ -211,3 +211,261 @@ def compute_fn(y_preds, y_targets): assert torch.equal(em._targets[0].cpu(), output1[1].cpu()) assert torch.equal(em._targets[1].cpu(), output2[1].cpu()) assert em.compute() == 0.0 + + +def test_epoch_metric_compute_fn_tensor_output(): + """Test EpochMetric with compute_fn returning a tensor.""" + + def compute_fn(y_preds, y_targets): + return torch.mean(((y_preds - y_targets.type_as(y_preds)) ** 2), dim=0) + + em = EpochMetric(compute_fn) + em.reset() + output1 = (torch.rand(4, 3), torch.randint(0, 2, size=(4, 3), dtype=torch.long)) + em.update(output1) + output2 = (torch.rand(4, 3), torch.randint(0, 2, size=(4, 3), dtype=torch.long)) + em.update(output2) + + result = em.compute() + assert isinstance(result, torch.Tensor) + assert result.shape == (3,) + + preds = torch.cat([output1[0], output2[0]], dim=0) + targets = torch.cat([output1[1], output2[1]], dim=0) + expected = compute_fn(preds, targets) + assert torch.allclose(result, expected) + + +def test_epoch_metric_compute_fn_tuple_output(): + """Test EpochMetric with compute_fn returning a tuple of tensors.""" + + def compute_fn(y_preds, y_targets): + mse = torch.mean(((y_preds - y_targets.type_as(y_preds)) ** 2)) + mae = torch.mean(torch.abs(y_preds - y_targets.type_as(y_preds))) + return (mse, mae) + + em = EpochMetric(compute_fn) + em.reset() + output1 = (torch.rand(4, 3), torch.randint(0, 2, size=(4, 3), dtype=torch.long)) + em.update(output1) + output2 = (torch.rand(4, 3), torch.randint(0, 2, size=(4, 3), dtype=torch.long)) + em.update(output2) + + result = em.compute() + assert isinstance(result, tuple) + assert len(result) == 2 + + preds = torch.cat([output1[0], output2[0]], dim=0) + targets = torch.cat([output1[1], output2[1]], dim=0) + expected = compute_fn(preds, targets) + assert torch.allclose(result[0], expected[0]) + assert torch.allclose(result[1], expected[1]) + + +def test_epoch_metric_compute_fn_invalid_output(): + """Test EpochMetric raises TypeError for unsupported compute_fn output.""" + + def compute_fn(y_preds, y_targets): + return "invalid_output" + + em = EpochMetric(compute_fn, check_compute_fn=False) + em.reset() + output1 = (torch.rand(4, 3), torch.randint(0, 2, size=(4, 3), dtype=torch.long)) + em.update(output1) + output2 = (torch.rand(4, 3), torch.randint(0, 2, size=(4, 3), dtype=torch.long)) + em.update(output2) + + with pytest.raises(TypeError, match=r"compute_fn output type"): + em.compute() + + +def test_epoch_metric_compute_fn_list_output(): + """Test EpochMetric with compute_fn returning a list of tensors.""" + + def compute_fn(y_preds, y_targets): + mse = torch.mean(((y_preds - y_targets.type_as(y_preds)) ** 2)) + mae = torch.mean(torch.abs(y_preds - y_targets.type_as(y_preds))) + return [mse, mae] + + em = EpochMetric(compute_fn) + em.reset() + output1 = (torch.rand(4, 3), torch.randint(0, 2, size=(4, 3), dtype=torch.long)) + em.update(output1) + output2 = (torch.rand(4, 3), torch.randint(0, 2, size=(4, 3), dtype=torch.long)) + em.update(output2) + + result = em.compute() + assert isinstance(result, list) + assert len(result) == 2 + + preds = torch.cat([output1[0], output2[0]], dim=0) + targets = torch.cat([output1[1], output2[1]], dim=0) + expected = compute_fn(preds, targets) + assert torch.allclose(result[0], expected[0]) + assert torch.allclose(result[1], expected[1]) + + +def test_epoch_metric_compute_fn_dict_output(): + """Test EpochMetric with compute_fn returning a dict of tensors.""" + + def compute_fn(y_preds, y_targets): + return { + "mse": torch.mean(((y_preds - y_targets.type_as(y_preds)) ** 2)), + "mae": torch.mean(torch.abs(y_preds - y_targets.type_as(y_preds))), + } + + em = EpochMetric(compute_fn) + em.reset() + output1 = (torch.rand(4, 3), torch.randint(0, 2, size=(4, 3), dtype=torch.long)) + em.update(output1) + output2 = (torch.rand(4, 3), torch.randint(0, 2, size=(4, 3), dtype=torch.long)) + em.update(output2) + + result = em.compute() + assert isinstance(result, dict) + assert "mse" in result + assert "mae" in result + + preds = torch.cat([output1[0], output2[0]], dim=0) + targets = torch.cat([output1[1], output2[1]], dim=0) + expected = compute_fn(preds, targets) + assert torch.allclose(result["mse"], expected["mse"]) + assert torch.allclose(result["mae"], expected["mae"]) + + +def test_epoch_metric_nested_invalid_output_raises(): + """Test EpochMetric raises TypeError for container with invalid nested type.""" + + def compute_fn(y_preds, y_targets): + return [torch.tensor(1.0), "not-a-number"] + + em = EpochMetric(compute_fn, check_compute_fn=False) + em.reset() + em.update((torch.rand(4, 3), torch.randint(0, 2, size=(4, 3), dtype=torch.long))) + em.update((torch.rand(4, 3), torch.randint(0, 2, size=(4, 3), dtype=torch.long))) + + with pytest.raises(TypeError, match=r"compute_fn output type .* is not supported"): + em.compute() + + +def test_epoch_metric_mapping_non_str_key_raises(): + """Test EpochMetric raises TypeError for mapping with non-string keys.""" + + def compute_fn(y_preds, y_targets): + return {0: torch.tensor(1.0)} + + em = EpochMetric(compute_fn, check_compute_fn=False) + em.reset() + em.update((torch.rand(4, 3), torch.randint(0, 2, size=(4, 3), dtype=torch.long))) + em.update((torch.rand(4, 3), torch.randint(0, 2, size=(4, 3), dtype=torch.long))) + + with pytest.raises(TypeError, match=r"mapping keys should be str"): + em.compute() + + +def test_distrib_container_outputs(distributed): + """Test EpochMetric with container outputs in distributed setting.""" + device = idist.device() if idist.device().type != "xla" else "cpu" + rank = idist.get_rank() + torch.manual_seed(40 + rank) + + n_iters = 3 + batch_size = 2 + n_classes = 7 + + y_true = torch.randint(0, n_classes, size=(n_iters * batch_size,), device=device) + y_preds = torch.rand(n_iters * batch_size, n_classes, device=device) + + def update(engine, i): + return ( + y_preds[i * batch_size : (i + 1) * batch_size, :], + y_true[i * batch_size : (i + 1) * batch_size], + ) + + engine = Engine(update) + + # Test tuple output + def tuple_fn(preds, targets): + return (torch.tensor(1.0), torch.tensor(2.0)) + + ep_metric = EpochMetric(tuple_fn, check_compute_fn=False, device=device) + ep_metric.attach(engine, "tup") + + # Test dict output + def dict_fn(preds, targets): + return {"a": torch.tensor(3.0), "b": torch.tensor(4.0)} + + ep_metric2 = EpochMetric(dict_fn, check_compute_fn=False, device=device) + ep_metric2.attach(engine, "dct") + + engine.run(data=list(range(n_iters)), max_epochs=1) + + # Verify tuple + tup = engine.state.metrics["tup"] + assert isinstance(tup, tuple) + assert len(tup) == 2 + assert torch.allclose(tup[0], torch.tensor(1.0)) + assert torch.allclose(tup[1], torch.tensor(2.0)) + + # Verify dict + dct = engine.state.metrics["dct"] + assert isinstance(dct, dict) + assert torch.allclose(dct["a"], torch.tensor(3.0)) + assert torch.allclose(dct["b"], torch.tensor(4.0)) + + +def test_distrib_invalid_output_raises_on_all_ranks(distributed): + """Regression test: an unsupported compute_fn output must raise TypeError on every + rank, not just rank 0. compute_fn only runs on rank 0, so if that rank validated the + output and raised before the broadcast collective started, rank 0 would exit while + other ranks hung forever waiting on a broadcast rank 0 never issued. + """ + device = idist.device() if idist.device().type != "xla" else "cpu" + + def compute_fn(y_preds, y_targets): + return "not-a-supported-type" + + em = EpochMetric(compute_fn, check_compute_fn=False, device=device) + em.reset() + em.update((torch.rand(4, 3, device=device), torch.randint(0, 2, size=(4, 3), device=device, dtype=torch.long))) + + with pytest.raises(TypeError, match=r"compute_fn output type.*is not supported"): + em.compute() + + +def test_distrib_mapping_non_str_key_raises_on_all_ranks(distributed): + """Regression test: a mapping output with a non-str key must raise TypeError on every + rank, for the same reason as test_distrib_invalid_output_raises_on_all_ranks above. + """ + device = idist.device() if idist.device().type != "xla" else "cpu" + + def compute_fn(y_preds, y_targets): + return {0: torch.tensor(1.0, device=device)} + + em = EpochMetric(compute_fn, check_compute_fn=False, device=device) + em.reset() + em.update((torch.rand(4, 3, device=device), torch.randint(0, 2, size=(4, 3), device=device, dtype=torch.long))) + + with pytest.raises(TypeError, match=r"mapping keys should be str"): + em.compute() + + +def test_distrib_nested_container_outputs(distributed): + """Test EpochMetric broadcasts nested containers (e.g. a dict holding a tuple and an + int) correctly, not just a single level of tuple/list/dict. + """ + device = idist.device() if idist.device().type != "xla" else "cpu" + + def compute_fn(y_preds, y_targets): + return {"scores": (torch.tensor(1.0, device=device), torch.tensor(2.0, device=device)), "count": 3} + + em = EpochMetric(compute_fn, check_compute_fn=False, device=device) + em.reset() + em.update((torch.rand(4, 3, device=device), torch.randint(0, 2, size=(4, 3), device=device, dtype=torch.long))) + + result = em.compute() + assert isinstance(result, dict) + assert isinstance(result["scores"], tuple) + assert torch.allclose(result["scores"][0].cpu(), torch.tensor(1.0)) + assert torch.allclose(result["scores"][1].cpu(), torch.tensor(2.0)) + assert result["count"] == 3