remove sklearn dependency from cohenkappa score calculation logic and applied custom calculation and updated tests (#3731)

Fixes #3701

  ## Description:
Remove scikit-learn dependency from `CohenKappa` by implementing a
native PyTorch version.

- Replaced `sklearn.metrics.cohen_kappa_score` with a pure PyTorch
implementation
- Removed forced GPU→CPU transfer (`.cpu().numpy()`) — metric now runs
fully on the configured device
- Built confusion matrix with `torch.bincount` (single vectorised
kernel) instead of a Python loop
  - Used `float64` throughout to match sklearn's numerical precision
- Added explicit multilabel input validation (previously handled
implicitly by sklearn
  
Check list:

- [x] New tests are added (if a new feature is added)
- [x] New doc strings: description and/or example code are in RST format
- [x] Documentation is updated (if required)

---------

Co-authored-by: vfdev <vfdev.5@gmail.com>

Author: avishkarsonni

ignite/metrics/cohen_kappa.py

@@ -2,15 +2,156 @@
 from typing import Literal
 
 import torch
+import torch.nn.functional as F
 
+from ignite.exceptions import NotComputableError
+from ignite.metrics.confusion_matrix import ConfusionMatrix
 from ignite.metrics.epoch_metric import EpochMetric
+from ignite.metrics.metric import Metric, reinit__is_reduced
 
 
-class CohenKappa(EpochMetric):
-    """Compute different types of Cohen's Kappa: Non-Wieghted, Linear, Quadratic.
-    Accumulating predictions and the ground-truth during an epoch and applying
-    `sklearn.metrics.cohen_kappa_score <https://scikit-learn.org/stable/modules/
-    generated/sklearn.metrics.cohen_kappa_score.html>`_ .
+def _kappa_from_conf(conf: torch.Tensor, weights: Literal["linear", "quadratic"] | None) -> float:
+    n = conf.sum()
+    if n == 0:
+        raise NotComputableError("CohenKappa cannot be computed on an empty confusion matrix (n == 0).")
+
+    n_classes = conf.shape[0]
+
+    if weights is None:
+        p_o = conf.trace() / n
+        row = conf.sum(dim=1)
+        col = conf.sum(dim=0)
+        p_e = (row * col).sum() / (n * n)
+    else:
+        idx = torch.arange(n_classes, device=conf.device)
+        if weights == "linear":
+            w = torch.abs(idx.unsqueeze(0) - idx.unsqueeze(1)).to(dtype=conf.dtype)
+        else:
+            w = ((idx.unsqueeze(0) - idx.unsqueeze(1)) ** 2).to(dtype=conf.dtype)
+
+        w = w / w.max()
+        p_o = 1 - (w * conf).sum() / n
+        row = conf.sum(dim=1)
+        col = conf.sum(dim=0)
+        expected = row.unsqueeze(1) * col.unsqueeze(0) / n
+        p_e = 1 - (w * expected).sum() / n
+
+    epsilon = 1e-9
+    return ((p_o - p_e) / (1 - p_e).clamp(min=epsilon)).item()
+
+
+def _cohen_kappa_score(
+    y_pred: torch.Tensor,
+    y: torch.Tensor,
+    weights: Literal["linear", "quadratic"] | None,
+    double_dtype: torch.dtype,
+) -> float:
+    if y_pred.ndim > 1 or y.ndim > 1:
+        raise ValueError("multilabel-indicator is not supported")
+
+    num_classes = int(max(y_pred.max().item(), y.max().item())) + 1
+
+    # Build the confusion matrix locally with plain tensor ops. We must NOT use the
+    # ``ConfusionMatrix`` metric here: its ``compute()`` is wrapped with ``sync_all_reduce``
+    # and this function runs on rank 0 only (see ``EpochMetric.compute``), so a collective
+    # would deadlock the other ranks.
+    y, y_pred = y.long(), y_pred.long()
+    indices = num_classes * y + y_pred
+    conf = torch.bincount(indices, minlength=num_classes**2)
+    conf = conf.reshape(num_classes, num_classes).to(dtype=double_dtype)
+
+    return _kappa_from_conf(conf, weights)
+
+
+class _CohenKappaEpochMetric(EpochMetric):
+    """CohenKappa backed by EpochMetric — infers num_classes dynamically from data."""
+
+    def __init__(
+        self,
+        weights: Literal["linear", "quadratic"] | None,
+        output_transform: Callable,
+        device: str | torch.device,
+        skip_unrolling: bool,
+        check_compute_fn: bool,
+    ):
+        super().__init__(
+            # ``self._double_dtype`` (set in Metric.__init__, float32 on MPS / float64 otherwise)
+            # is resolved lazily at compute time.
+            compute_fn=lambda y_pred, y: _cohen_kappa_score(y_pred, y, weights, self._double_dtype),
+            output_transform=output_transform,
+            check_compute_fn=check_compute_fn,
+            device=device,
+            skip_unrolling=skip_unrolling,
+        )
+
+    @reinit__is_reduced
+    def update(self, output: tuple[torch.Tensor, torch.Tensor]) -> None:
+        y_pred, y = output[0].detach(), output[1].detach()
+
+        if y_pred.ndim == 2 and y_pred.shape[1] == 1:
+            y_pred = y_pred.squeeze(dim=-1)
+        if y.ndim == 2 and y.shape[1] == 1:
+            y = y.squeeze(dim=-1)
+
+        super().update((y_pred, y))
+
+
+class _CohenKappaConfusionMatrix(Metric):
+    """CohenKappa backed by ConfusionMatrix — requires num_classes at construction time.
+    Accumulates a running confusion matrix; no raw tensor buffering.
+    """
+
+    _state_dict_all_req_keys = ("_cm",)
+
+    def __init__(
+        self,
+        num_classes: int,
+        weights: Literal["linear", "quadratic"] | None,
+        output_transform: Callable,
+        device: str | torch.device,
+        skip_unrolling: bool,
+    ):
+        self._weights = weights
+        self._cm = ConfusionMatrix(
+            num_classes=num_classes,
+            output_transform=output_transform,
+            device=device,
+            skip_unrolling=skip_unrolling,
+        )
+        super().__init__(output_transform=output_transform, device=device, skip_unrolling=skip_unrolling)
+
+    @reinit__is_reduced
+    def reset(self) -> None:
+        self._cm.reset()
+
+    @reinit__is_reduced
+    def update(self, output: tuple[torch.Tensor, torch.Tensor]) -> None:
+        y_pred, y = output[0].detach(), output[1].detach()
+
+        if y_pred.ndim == 2 and y_pred.shape[1] == 1:
+            y_pred = y_pred.squeeze(dim=-1)
+        if y.ndim == 2 and y.shape[1] == 1:
+            y = y.squeeze(dim=-1)
+
+        if y_pred.ndim > 1 or y.ndim > 1:
+            raise ValueError("multilabel-indicator is not supported")
+
+        num_classes = self._cm.num_classes
+        y_pred_oh = F.one_hot(y_pred.long(), num_classes).float().to(self._device)
+        self._cm.update((y_pred_oh, y.long().to(self._device)))
+
+    def compute(self) -> float:
+        conf = self._cm.compute().to(dtype=self._double_dtype)
+        return _kappa_from_conf(conf, self._weights)
+
+
+class CohenKappa(Metric):
+    """Compute different types of Cohen's Kappa: Non-Weighted, Linear, Quadratic.
+
+    When ``num_classes`` is provided, accumulates a running confusion matrix via
+    :class:`~ignite.metrics.confusion_matrix.ConfusionMatrix` (memory-efficient, no raw tensor buffering).
+    When ``num_classes`` is ``None`` (default), buffers predictions and targets via
+    :class:`~ignite.metrics.EpochMetric` and infers the number of classes from data.
 
     Args:
         output_transform: a callable that is used to transform the
@@ -19,19 +160,20 @@ class CohenKappa(EpochMetric):
             you want to compute the metric with respect to one of the outputs.
         weights: a string is used to define the type of Cohen's Kappa whether Non-Weighted or Linear
             or Quadratic. Default, None.
-        check_compute_fn: Default False. If True, `cohen_kappa_score
-            <https://scikit-learn.org/stable/modules/generated/sklearn.metrics.cohen_kappa_score.html>`_
-            is run on the first batch of data to ensure there are
-            no issues. User will be warned in case there are any issues computing the function.
+        check_compute_fn: Default False. If True, the compute function is run on the first batch
+            of data to ensure there are no issues. User will be warned in case there are any issues
+            computing the function.
         device: optional device specification for internal storage.
         skip_unrolling: specifies whether output should be unrolled before being fed to update method. Should be
             true for multi-output model, for example, if ``y_pred`` contains multi-output as ``(y_pred_a, y_pred_b)``
             Alternatively, ``output_transform`` can be used to handle this.
+        num_classes: number of classes. If provided, uses a running confusion matrix
+            (memory-efficient). If ``None``, infers from data at compute time (backward-compatible default).
 
     Examples:
         To use with ``Engine`` and ``process_function``, simply attach the metric instance to the engine.
         The output of the engine's ``process_function`` needs to be in the format of
-        ``(y_pred, y)`` or ``{'y_pred': y_pred, 'y': y, ...}``. If not, ``output_tranform`` can be added
+        ``(y_pred, y)`` or ``{'y_pred': y_pred, 'y': y, ...}``. If not, ``output_transform`` can be added
         to the metric to transform the output into the form expected by the metric.
 
         .. include:: defaults.rst
@@ -52,6 +194,10 @@ class CohenKappa(EpochMetric):
 
     .. versionchanged:: 0.5.1
         ``skip_unrolling`` argument is added.
+
+    .. versionchanged:: 0.6.0
+        Replaced scikit-learn dependency with a native PyTorch implementation.
+        Added ``num_classes`` argument; routes to a running-confusion-matrix backend when provided.
     """
 
     def __init__(
@@ -61,28 +207,39 @@ def __init__(
         check_compute_fn: bool = False,
         device: str | torch.device = torch.device("cpu"),
         skip_unrolling: bool = False,
+        num_classes: int | None = None,
     ):
-        try:
-            from sklearn.metrics import cohen_kappa_score  # noqa: F401
-        except ImportError:
-            raise ModuleNotFoundError("This contrib module requires scikit-learn to be installed.")
         if weights not in (None, "linear", "quadratic"):
             raise ValueError("Kappa Weighting type must be None or linear or quadratic.")
 
-        # initialize weights
         self.weights: Literal["linear", "quadratic"] | None = weights
 
-        super().__init__(
-            self._cohen_kappa_score,
-            output_transform=output_transform,
-            check_compute_fn=check_compute_fn,
-            device=device,
-            skip_unrolling=skip_unrolling,
-        )
-
-    def _cohen_kappa_score(self, y_targets: torch.Tensor, y_preds: torch.Tensor) -> float:
-        from sklearn.metrics import cohen_kappa_score
-
-        y_true = y_targets.cpu().numpy()
-        y_pred = y_preds.cpu().numpy()
-        return cohen_kappa_score(y_true, y_pred, weights=self.weights)
+        if num_classes is not None:
+            self._impl: Metric = _CohenKappaConfusionMatrix(
+                num_classes=num_classes,
+                weights=weights,
+                output_transform=output_transform,
+                device=device,
+                skip_unrolling=skip_unrolling,
+            )
+        else:
+            self._impl = _CohenKappaEpochMetric(
+                weights=weights,
+                output_transform=output_transform,
+                device=device,
+                skip_unrolling=skip_unrolling,
+                check_compute_fn=check_compute_fn,
+            )
+
+        super().__init__(output_transform=output_transform, device=device, skip_unrolling=skip_unrolling)
+
+    @reinit__is_reduced
+    def reset(self) -> None:
+        self._impl.reset()
+
+    @reinit__is_reduced
+    def update(self, output: tuple[torch.Tensor, torch.Tensor]) -> None:
+        self._impl.update(output)
+
+    def compute(self) -> float:
+        return self._impl.compute()

ignite/metrics/epoch_metric.py

@@ -144,7 +144,7 @@ def update(self, output: tuple[torch.Tensor, torch.Tensor]) -> None:
 
     def compute(self) -> float:
         if len(self._predictions) < 1 or len(self._targets) < 1:
-            raise NotComputableError("EpochMetric must have at least one example before it can be computed.")
+            raise NotComputableError(f"{type(self).__name__} must have at least one example before it can be computed.")
 
         if self._result is None:
             _prediction_tensor = torch.cat(self._predictions, dim=0)

tests/ignite/metrics/regression/test_median_absolute_error.py

@@ -13,7 +13,7 @@
 def test_zero_sample():
     m = MedianAbsoluteError()
     with pytest.raises(
-        NotComputableError, match=r"EpochMetric must have at least one example before it can be computed"
+        NotComputableError, match=r"MedianAbsoluteError must have at least one example before it can be computed"
     ):
         m.compute()
 

tests/ignite/metrics/regression/test_median_absolute_percentage_error.py

@@ -13,7 +13,8 @@
 def test_zero_sample():
     m = MedianAbsolutePercentageError()
     with pytest.raises(
-        NotComputableError, match=r"EpochMetric must have at least one example before it can be computed"
+        NotComputableError,
+        match=r"MedianAbsolutePercentageError must have at least one example before it can be computed",
     ):
         m.compute()
 

tests/ignite/metrics/regression/test_median_relative_absolute_error.py

@@ -13,7 +13,8 @@
 def test_zero_sample():
     m = MedianRelativeAbsoluteError()
     with pytest.raises(
-        NotComputableError, match=r"EpochMetric must have at least one example before it can be computed"
+        NotComputableError,
+        match=r"MedianRelativeAbsoluteError must have at least one example before it can be computed",
     ):
         m.compute()
 

tests/ignite/metrics/test_average_precision.py

@@ -29,7 +29,7 @@ def test_no_update():
     ap = AveragePrecision()
 
     with pytest.raises(
-        NotComputableError, match=r"EpochMetric must have at least one example before it can be computed"
+        NotComputableError, match=r"AveragePrecision must have at least one example before it can be computed"
     ):
         ap.compute()