implement registry based topk wrapper

Author: rwtarpit

ignite/metrics/precision.py

@@ -88,6 +88,10 @@ def _prepare_output(self, output: Sequence[torch.Tensor]) -> Sequence[torch.Tens
 
         return y_pred, y, correct
 
+    @classmethod
+    def _topk_transform(cls, output: Sequence[torch.tensor], top_k: int) -> Sequence[torch.tensor]:
+        return output
+
     @reinit__is_reduced
     def reset(self) -> None:
         """
@@ -410,8 +414,9 @@ def update(self, output: Sequence[torch.Tensor]) -> None:
                 For binary and multilabel data, both y and y_pred should consist of 0's and 1's, but for multiclass
                 data, y_pred and y should consist of probabilities and integers respectively.
         """
-        self._check_shape(output)
-        self._check_type(output)
+        if not getattr(self, "_skip_checks", False):
+            self._check_shape(output)
+            self._check_type(output)
         y_pred, y, correct = self._prepare_output(output)
 
         if self._average == "samples":

ignite/metrics/top_k.py

@@ -0,0 +1,81 @@
+import torch
+from typing import Sequence
+
+from ignite.metrics import Metric
+from ignite.metrics.precision import _BasePrecisionRecall
+
+
+class TopK(Metric):
+    """https://github.com/open-mmlab/mmengine/blob/main/mmengine/registry/registry.py
+        https://github.com/open-mmlab/mmdetection/tree/main
+
+    the idea is to maintain tok_k transforms here for metrics.
+    and user will have to pass output_transform to TopK instead of base_metric
+    and the output transform must only unpack output into y_pred, y, and do no other transformations like binarisation
+    """
+
+    _output_transform_registry = {}
+
+    @classmethod
+    def register(cls, metric_type, k_transform):
+        cls._output_transform_registry[metric_type] = k_transform
+
+    def __init__(
+        self,
+        base_metric: Metric,
+        top_k: int | list[int],
+        output_transform=lambda x: x,
+        device: str | torch.device = torch.device("cpu"),
+        skip_unrolling: bool = False,
+    ):
+        transform = None
+        for metric_type, k_transform in self._output_transform_registry.items():
+            if isinstance(base_metric, metric_type):
+                transform = k_transform
+
+        if transform is None:
+            raise ValueError(f"{type(base_metric).__name__} does not support TopK.")
+
+        self._transform = transform
+        self._base_metric = base_metric
+        self._ks = sorted(top_k) if isinstance(top_k, list) else [top_k]
+        super().__init__(output_transform=output_transform, device=device, skip_unrolling=skip_unrolling)
+
+    def reset(self):
+        self._base_metric.reset()
+        self._states = {k: self._base_metric.state_dict() for k in self._ks}
+
+    def update(self, output):
+        self._base_metric._check_shape(output)
+        self._base_metric._check_type(output)
+        self._base_metric._skip_checks = True
+
+        for k in self._ks:
+            # restore state for this k
+            self._base_metric.load_state_dict(self._states[k])
+
+            k_output = self._transform(output, k)
+            self._base_metric.update(k_output)
+
+            # save state for this k
+            self._states[k] = self._base_metric.state_dict()
+
+        self._base_metric._skip_checks = False
+
+    def compute(self) -> list:
+        results = []
+        for k in self._ks:
+            self._base_metric.load_state_dict(self._states[k])
+            results.append(self._base_metric.compute())
+        return results
+
+
+def _precision_recall_topk_transform(output: Sequence[torch.Tensor], k: int):
+    y_pred, y = output[0], output[1]
+    _, top_indices = torch.topk(y_pred, k=k, dim=-1)
+    masked = torch.zeros_like(y_pred)
+    masked.scatter_(-1, top_indices, 1.0)
+    return (masked, y)
+
+
+TopK.register(_BasePrecisionRecall, _precision_recall_topk_transform)