diff --git a/ignite/metrics/epoch_metric.py b/ignite/metrics/epoch_metric.py index 672f843d1697..9008826ad344 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 cast, Union import torch @@ -10,6 +10,9 @@ __all__ = ["EpochMetric"] +# Supported return types for ``EpochMetric``'s ``compute_fn``. +EpochMetricOutput = Union[int, float, torch.Tensor, Sequence, Mapping] + class EpochMetric(Metric): """Class for metrics that should be computed on the entire output history of a model. @@ -30,7 +33,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. An unsupported return type raises a ``TypeError``. + Note: in distributed configuration (``world_size > 1``), only scalar and + ``torch.Tensor`` outputs are broadcast across processes; tuple/list/mapping outputs + are supported only when ``world_size == 1``. 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 +102,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 +151,27 @@ 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 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 +185,48 @@ 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)) + # All ranks must take the same path through the collective calls below, otherwise + # they would deadlock on mismatched broadcasts. Only rank 0 holds the real result, + # so it classifies the result and shares a status code with every rank *before* + # broadcasting the result itself. Type/validation problems are surfaced through the + # same mechanism so that every rank raises the same exception. + _BROADCASTABLE, _UNSUPPORTED_CONTAINER, _UNSUPPORTED_TYPE = 0, 1, 2 + status = _BROADCASTABLE + if idist.get_rank() == 0 and not isinstance(result, (int, float, torch.Tensor)): + try: + self._check_output_type(result) + status = _UNSUPPORTED_CONTAINER + except TypeError: + status = _UNSUPPORTED_TYPE + status = int(idist.broadcast(status, src=0)) + + if status == _UNSUPPORTED_TYPE: + # Every rank raises the same error; no result broadcast is attempted. + 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." + ) + if status == _UNSUPPORTED_CONTAINER: + # Every rank raises the same error; no result broadcast is attempted. + raise NotImplementedError( + "Distributed broadcast of tuple/list/mapping compute_fn outputs is not supported yet. " + "Such outputs are currently supported only in non-distributed (world_size == 1) " + "configuration." + ) + + # status == _BROADCASTABLE: every rank performs the matching result broadcast. + result = cast(EpochMetricOutput, idist.broadcast(result, src=0, safe_mode=True)) + else: + # Single process: validate directly and surface unsupported types as TypeError. + self._check_output_type(result) + + self._result = result return self._result diff --git a/tests/ignite/metrics/test_epoch_metric.py b/tests/ignite/metrics/test_epoch_metric.py index 5bbb2e2307cc..6b5e03f87ac9 100644 --- a/tests/ignite/metrics/test_epoch_metric.py +++ b/tests/ignite/metrics/test_epoch_metric.py @@ -211,3 +211,234 @@ 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_scalar_tensor_output(available_device): + # compute_fn returns a 0-dim (scalar) tensor instead of a python float + def compute_fn(y_preds, y_targets): + return torch.mean(y_preds - y_targets.type_as(y_preds)) + + em = EpochMetric(compute_fn, device=available_device) + + 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) + + preds = torch.cat([output1[0], output2[0]], dim=0) + targets = torch.cat([output1[1], output2[1]], dim=0) + expected = compute_fn(preds, targets) + + result = em.compute() + assert isinstance(result, torch.Tensor) + assert result.ndim == 0 + assert torch.allclose(result.cpu(), expected.cpu()) + + +def test_epoch_metric_vector_tensor_output(available_device): + # compute_fn returns a 1-dim (vector) tensor, one value per target column + def compute_fn(y_preds, y_targets): + return y_preds.mean(dim=0) + + em = EpochMetric(compute_fn, device=available_device) + + 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) + + preds = torch.cat([output1[0], output2[0]], dim=0) + targets = torch.cat([output1[1], output2[1]], dim=0) + expected = compute_fn(preds, targets) + + result = em.compute() + assert isinstance(result, torch.Tensor) + assert result.shape == (3,) + assert torch.allclose(result.cpu(), expected.cpu()) + + +def test_epoch_metric_tuple_of_tensors_output(available_device): + # compute_fn returns a tuple of tensors + def compute_fn(y_preds, y_targets): + return y_preds.mean(dim=0), y_preds.sum(dim=0) + + em = EpochMetric(compute_fn, device=available_device) + + 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) + + preds = torch.cat([output1[0], output2[0]], dim=0) + targets = torch.cat([output1[1], output2[1]], dim=0) + expected = compute_fn(preds, targets) + + result = em.compute() + assert isinstance(result, tuple) + assert len(result) == 2 + for r, e in zip(result, expected): + assert torch.allclose(r.cpu(), e.cpu()) + + +def test_epoch_metric_list_of_tensors_output(available_device): + # compute_fn returns a list of tensors + def compute_fn(y_preds, y_targets): + return [y_preds.mean(dim=0), y_preds.sum(dim=0)] + + em = EpochMetric(compute_fn, device=available_device) + + 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) + + preds = torch.cat([output1[0], output2[0]], dim=0) + targets = torch.cat([output1[1], output2[1]], dim=0) + expected = compute_fn(preds, targets) + + result = em.compute() + assert isinstance(result, list) + assert len(result) == 2 + for r, e in zip(result, expected): + assert torch.allclose(r.cpu(), e.cpu()) + + +def test_epoch_metric_mapping_of_tensors_output(available_device): + # compute_fn returns a dict mapping str -> tensor + def compute_fn(y_preds, y_targets): + return {"mean": y_preds.mean(dim=0), "sum": y_preds.sum(dim=0)} + + em = EpochMetric(compute_fn, device=available_device) + + 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) + + preds = torch.cat([output1[0], output2[0]], dim=0) + targets = torch.cat([output1[1], output2[1]], dim=0) + expected = compute_fn(preds, targets) + + result = em.compute() + assert isinstance(result, dict) + assert set(result.keys()) == {"mean", "sum"} + for key in expected: + assert torch.allclose(result[key].cpu(), expected[key].cpu()) + + +def test_epoch_metric_unsupported_output_type_raises(available_device): + # An unsupported output type (str) should raise a clear TypeError. + def compute_fn(y_preds, y_targets): + return "not-a-number" + + em = EpochMetric(compute_fn, check_compute_fn=False, device=available_device) + + 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_nested_invalid_output_raises(available_device): + # A container holding an unsupported value (str) should be rejected by the recursive check. + def compute_fn(y_preds, y_targets): + return [torch.tensor(1.0), "not-a-number"] + + em = EpochMetric(compute_fn, check_compute_fn=False, device=available_device) + + 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(available_device): + # A mapping with a non-string key is not a supported output. + def compute_fn(y_preds, y_targets): + return {0: torch.tensor(1.0)} + + em = EpochMetric(compute_fn, check_compute_fn=False, device=available_device) + + 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_output_tensor(distributed): + # A non-scalar tensor output should broadcast successfully and be identical on all ranks. + device = idist.device() if idist.device().type != "xla" else "cpu" + torch.manual_seed(40 + idist.get_rank()) + + def compute_fn(y_preds, y_targets): + return y_preds.mean(dim=0) + + 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), dtype=torch.long, device=device))) + em.update((torch.rand(4, 3, device=device), torch.randint(0, 2, size=(4, 3), dtype=torch.long, device=device))) + + result = em.compute() + assert isinstance(result, torch.Tensor) + assert result.shape == (3,) + # All ranks must receive the same broadcasted values. + gathered = idist.all_gather(result.unsqueeze(0)) + assert torch.allclose(gathered[0], gathered[-1]) + + +@pytest.mark.parametrize( + "compute_fn", + [ + lambda y_preds, y_targets: (y_preds.mean(dim=0), y_preds.sum(dim=0)), + lambda y_preds, y_targets: [y_preds.mean(dim=0), y_preds.sum(dim=0)], + lambda y_preds, y_targets: {"mean": y_preds.mean(dim=0), "sum": y_preds.sum(dim=0)}, + ], + ids=["tuple", "list", "dict"], +) +def test_distrib_output_container_raises(distributed, compute_fn): + # Documents the current conservative distributed behavior: tuple/list/mapping outputs are + # supported in single-process mode, but under world_size > 1 they intentionally fail + # symmetrically with NotImplementedError rather than entering mismatched collective calls. + # This is a safety/regression guarantee for this implementation, not the ideal long-term + # feature behavior (broadcasting containers across ranks is left for future work). + if idist.get_world_size() == 1: + pytest.skip("This test verifies world_size > 1 behavior; containers are supported in single process.") + + device = idist.device() if idist.device().type != "xla" else "cpu" + torch.manual_seed(40 + idist.get_rank()) + + 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), dtype=torch.long, device=device))) + em.update((torch.rand(4, 3, device=device), torch.randint(0, 2, size=(4, 3), dtype=torch.long, device=device))) + + with pytest.raises(NotImplementedError, match=r"Distributed broadcast of tuple/list/mapping"): + em.compute() + + +def test_distrib_output_unsupported_raises(distributed): + # An unsupported output type (str) should raise TypeError consistently on all ranks. + device = idist.device() if idist.device().type != "xla" else "cpu" + torch.manual_seed(40 + idist.get_rank()) + + def compute_fn(y_preds, y_targets): + return "not-a-number" + + 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), dtype=torch.long, device=device))) + em.update((torch.rand(4, 3, device=device), torch.randint(0, 2, size=(4, 3), dtype=torch.long, device=device))) + + with pytest.raises(TypeError, match=r"compute_fn output type.*is not supported"): + em.compute()