New metric: Cluster Accuracy

.github/workflows/ci-tests.yml

@@ -63,6 +63,7 @@ jobs:
       TOKENIZERS_PARALLELISM: false
       TEST_DIRS: ${{ needs.check-diff.outputs.test-dirs }}
       PIP_EXTRA_INDEX_URL: "--find-links https://download.pytorch.org/whl/cpu/torch_stable.html"
+      PIP_USE_FEATURE: "2020-resolver"
 
     # Timeout: https://stackoverflow.com/a/59076067/4521646
     # seems that macOS jobs take much more than orger OS

CHANGELOG.md

@@ -12,6 +12,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Added
 
+- Added `ClusterAccuracy` metric to cluster package ([#2777](https://github.com/Lightning-AI/torchmetrics/pull/2777))
+
+
 -
 
 

docs/source/clustering/cluster_accuracy.rst

@@ -0,0 +1,21 @@
+.. customcarditem::
+   :header: Cluster Accuracy
+   :image: https://pl-flash-data.s3.amazonaws.com/assets/thumbnails/default.svg
+   :tags: Clustering
+
+.. include:: ../links.rst
+
+################
+Cluster Accuracy
+################
+
+Module Interface
+________________
+
+.. autoclass:: torchmetrics.clustering.ClusterAccuracy
+    :exclude-members: update, compute
+
+Functional Interface
+____________________
+
+.. autofunction:: torchmetrics.functional.clustering.cluster_accuracy

docs/source/links.rst

@@ -177,5 +177,6 @@
 .. _Hausdorff Distance: https://en.wikipedia.org/wiki/Hausdorff_distance
 .. _averaging curve objects: https://scikit-learn.org/stable/auto_examples/model_selection/plot_roc.html
 .. _Procrustes Disparity: https://en.wikipedia.org/wiki/Procrustes_analysis
+.. _Cluster Accuracy: https://arxiv.org/abs/2206.07579
 .. _Log AUC: https://pubmed.ncbi.nlm.nih.gov/20735049/
 .. _Negative Predictive Value: https://en.wikipedia.org/wiki/Positive_and_negative_predictive_values

requirements/_devel.txt

@@ -5,19 +5,21 @@
 -r _tests.txt
 
 # add extra requirements
+-r audio.txt
+-r clustering.txt
+-r detection.txt
 -r image.txt
 -r text.txt
--r detection.txt
--r audio.txt
 -r multimodal.txt
 -r visual.txt
 
 # add extra testing
--r image_test.txt
--r text_test.txt
 -r audio_test.txt
--r detection_test.txt
 -r classification_test.txt
+-r clustering_test.txt
+-r detection_test.txt
+-r image_test.txt
 -r nominal_test.txt
--r segmentation_test.txt
 -r regression_test.txt
+-r segmentation_test.txt
+-r text_test.txt

requirements/_docs.txt

@@ -18,12 +18,13 @@ pydantic > 1.0.0, < 3.0.0
 
 # integrations
 -r _integrate.txt
--r visual.txt
 -r audio.txt
+-r clustering.txt
 -r detection.txt
 -r image.txt
 -r multimodal.txt
 -r text.txt
+-r visual.txt
 
 # Gallery extra requirements
 # --------------------------

requirements/clustering.txt

@@ -0,0 +1,4 @@
+# NOTE: the upper bound for the package version is only set for CI stability, and it is dropped while installing this package
+#  in case you want to preserve/enforce restrictions on the latest compatible version, add "strict" as an in-line comment
+
+torch_linear_assignment <0.0.2

requirements/clustering_test.txt

@@ -0,0 +1,4 @@
+# NOTE: the upper bound for the package version is only set for CI stability, and it is dropped while installing this package
+#  in case you want to preserve/enforce restrictions on the latest compatible version, add "strict" as an in-line comment
+
+aeon >= 0.7.0; python_version >"3.10"  # cluster accuracy

src/torchmetrics/clustering/__init__.py

@@ -14,6 +14,7 @@
 from torchmetrics.clustering.adjusted_mutual_info_score import AdjustedMutualInfoScore
 from torchmetrics.clustering.adjusted_rand_score import AdjustedRandScore
 from torchmetrics.clustering.calinski_harabasz_score import CalinskiHarabaszScore
+from torchmetrics.clustering.cluster_accuracy import ClusterAccuracy
 from torchmetrics.clustering.davies_bouldin_score import DaviesBouldinScore
 from torchmetrics.clustering.dunn_index import DunnIndex
 from torchmetrics.clustering.fowlkes_mallows_index import FowlkesMallowsIndex
@@ -30,6 +31,7 @@
     "AdjustedMutualInfoScore",
     "AdjustedRandScore",
     "CalinskiHarabaszScore",
+    "ClusterAccuracy",
     "CompletenessScore",
     "DaviesBouldinScore",
     "DunnIndex",

src/torchmetrics/clustering/adjusted_mutual_info_score.py

@@ -75,7 +75,6 @@ class AdjustedMutualInfoScore(MutualInfoScore):
     plot_upper_bound: float = 1.0
     preds: List[Tensor]
     target: List[Tensor]
-    contingency: Tensor
 
     def __init__(
         self, average_method: Literal["min", "geometric", "arithmetic", "max"] = "arithmetic", **kwargs: Any

src/torchmetrics/clustering/cluster_accuracy.py

@@ -0,0 +1,148 @@
+# Copyright The Lightning team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from typing import Any, Optional, Sequence, Union
+
+import torch
+from torch import Tensor
+
+from torchmetrics.functional.classification import multiclass_confusion_matrix
+from torchmetrics.functional.clustering.cluster_accuracy import _cluster_accuracy_compute
+from torchmetrics.metric import Metric
+from torchmetrics.utilities.imports import (
+    _MATPLOTLIB_AVAILABLE,
+    _TORCH_LINEAR_ASSIGNMENT_AVAILABLE,
+)
+from torchmetrics.utilities.plot import _AX_TYPE, _PLOT_OUT_TYPE
+
+if not _MATPLOTLIB_AVAILABLE:
+    __doctest_skip__ = ["ClusterAccuracy.plot"]
+
+if not _TORCH_LINEAR_ASSIGNMENT_AVAILABLE:
+    __doctest_skip__ = ["ClusterAccuracy", "ClusterAccuracy.plot"]
+
+
+class ClusterAccuracy(Metric):
+    r"""Compute `Cluster Accuracy`_ between predicted and target clusters.
+
+    .. math::
+
+        \text{Cluster Accuracy} = \max_g \frac{1}{N} \sum_{n=1}^N \mathbb{1}_{g(p_n) = t_n}
+
+    Where :math:`g` is a function that maps predicted clusters :math:`p` to target clusters :math:`t`, :math:`N` is the
+    number of samples, :math:`p_n` is the predicted cluster for sample :math:`n`, :math:`t_n` is the target cluster for
+    sample :math:`n`, and :math:`\mathbb{1}` is the indicator function. The function :math:`g` is determined by solving
+    the linear sum assignment problem.
+
+    This clustering metric is an extrinsic measure, because it requires ground truth clustering labels, which may not
+    be available in practice since clustering in generally is used for unsupervised learning.
+
+    As input to ``forward`` and ``update`` the metric accepts the following input:
+
+    - ``preds`` (:class:`~torch.Tensor`): single integer tensor with shape ``(N,)`` with predicted cluster labels
+    - ``target`` (:class:`~torch.Tensor`): single integer tensor with shape ``(N,)`` with ground truth cluster labels
+
+    As output of ``forward`` and ``compute`` the metric returns the following output:
+
+    - ``acc_score`` (:class:`~torch.Tensor`): A tensor with the Cluster Accuracy score
+
+    Args:
+        num_classes: number of classes
+        kwargs: Additional keyword arguments, see :ref:`Metric kwargs` for more info.
+
+    Raises:
+        RuntimeError:
+            If ``torch_linear_assignment`` is not installed. To install, run ``pip install torchmetrics[clustering]``.
+        ValueError
+            If ``num_classes`` is not a positive integer
+
+    Example::
+        >>> import torch
+        >>> from torchmetrics.clustering import ClusterAccuracy
+        >>> preds = torch.tensor([0, 0, 1, 1])
+        >>> target = torch.tensor([1, 1, 0, 0])
+        >>> metric = ClusterAccuracy(num_classes=2)
+        >>> metric(preds, target)
+        tensor(1.)
+
+    """
+
+    is_differentiable: bool = False
+    higher_is_better: bool = True
+    full_state_update: bool = False
+    plot_lower_bound: float = 0.0
+    plot_upper_bound: float = 1.0
+    confmat: Tensor
+
+    def __init__(self, num_classes: int, **kwargs: Any) -> None:
+        super().__init__(**kwargs)
+        if not _TORCH_LINEAR_ASSIGNMENT_AVAILABLE:
+            raise RuntimeError(
+                "Missing `torch_linear_assignment`. Please install it with `pip install torchmetrics[clustering]`."
+            )
+
+        if not isinstance(num_classes, int) or num_classes <= 0:
+            raise ValueError("Argument `num_classes` should be a positive integer")
+        self.add_state(
+            "confmat", default=torch.zeros((num_classes, num_classes), dtype=torch.int64), dist_reduce_fx="sum"
+        )
+        self.num_classes = num_classes
+
+    def update(self, preds: Tensor, target: Tensor) -> None:
+        """Update the confusion matrix with the new predictions and targets."""
+        self.confmat += multiclass_confusion_matrix(preds, target, num_classes=self.num_classes)
+
+    def compute(self) -> Tensor:
+        """Computes the clustering accuracy."""
+        return _cluster_accuracy_compute(self.confmat)
+
+    def plot(self, val: Union[Tensor, Sequence[Tensor], None] = None, ax: Optional[_AX_TYPE] = None) -> _PLOT_OUT_TYPE:
+        """Plot a single or multiple values from the metric.
+
+        Args:
+            val: Either a single result from calling ``metric.forward`` or ``metric.compute``
+                or a list of these results. If no value is provided, will automatically call `metric.compute`
+                and plot that result.
+            ax: An matplotlib axis object. If provided will add plot to that axis
+
+        Returns:
+            Figure and Axes object
+
+        Raises:
+            ModuleNotFoundError:
+                If `matplotlib` is not installed
+
+        .. plot::
+            :scale: 75
+
+            >>> # Example plotting a single value
+            >>> import torch
+            >>> from torchmetrics.clustering import ClusterAccuracy
+            >>> metric = ClusterAccuracy(num_classes=4)
+            >>> metric.update(torch.randint(0, 4, (10,)), torch.randint(0, 4, (10,)))
+            >>> fig_, ax_ = metric.plot(metric.compute())
+
+        .. plot::
+            :scale: 75
+
+            >>> # Example plotting multiple values
+            >>> import torch
+            >>> from torchmetrics.clustering import ClusterAccuracy
+            >>> metric = ClusterAccuracy(num_classes=4)
+            >>> values = [ ]
+            >>> for _ in range(10):
+            ...     values.append(metric(torch.randint(0, 4, (10,)), torch.randint(0, 4, (10,))))
+            >>> fig_, ax_ = metric.plot(values)
+
+        """
+        return self._plot(val, ax)

src/torchmetrics/clustering/fowlkes_mallows_index.py

@@ -65,7 +65,6 @@ class FowlkesMallowsIndex(Metric):
     plot_upper_bound: float = 1.0
     preds: List[Tensor]
     target: List[Tensor]
-    contingency: Tensor
 
     def __init__(self, **kwargs: Any) -> None:
         super().__init__(**kwargs)

src/torchmetrics/clustering/mutual_info_score.py

@@ -70,7 +70,6 @@ class MutualInfoScore(Metric):
     plot_lower_bound: float = 0.0
     preds: List[Tensor]
     target: List[Tensor]
-    contingency: Tensor
 
     def __init__(self, **kwargs: Any) -> None:
         super().__init__(**kwargs)

src/torchmetrics/clustering/normalized_mutual_info_score.py

@@ -74,7 +74,6 @@ class NormalizedMutualInfoScore(MutualInfoScore):
     plot_upper_bound: float = 0.0
     preds: List[Tensor]
     target: List[Tensor]
-    contingency: Tensor
 
     def __init__(
         self, average_method: Literal["min", "geometric", "arithmetic", "max"] = "arithmetic", **kwargs: Any

src/torchmetrics/clustering/rand_score.py

@@ -68,7 +68,6 @@ class RandScore(Metric):
     plot_lower_bound: float = 0.0
     preds: List[Tensor]
     target: List[Tensor]
-    contingency: Tensor
 
     def __init__(self, **kwargs: Any) -> None:
         super().__init__(**kwargs)

src/torchmetrics/functional/clustering/__init__.py

@@ -14,6 +14,7 @@
 from torchmetrics.functional.clustering.adjusted_mutual_info_score import adjusted_mutual_info_score
 from torchmetrics.functional.clustering.adjusted_rand_score import adjusted_rand_score
 from torchmetrics.functional.clustering.calinski_harabasz_score import calinski_harabasz_score
+from torchmetrics.functional.clustering.cluster_accuracy import cluster_accuracy
 from torchmetrics.functional.clustering.davies_bouldin_score import davies_bouldin_score
 from torchmetrics.functional.clustering.dunn_index import dunn_index
 from torchmetrics.functional.clustering.fowlkes_mallows_index import fowlkes_mallows_index
@@ -30,6 +31,7 @@
     "adjusted_mutual_info_score",
     "adjusted_rand_score",
     "calinski_harabasz_score",
+    "cluster_accuracy",
     "completeness_score",
     "davies_bouldin_score",
     "dunn_index",

src/torchmetrics/functional/clustering/cluster_accuracy.py

@@ -0,0 +1,67 @@
+# Copyright The Lightning team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import torch
+from torch import Tensor
+
+from torchmetrics.functional.classification import multiclass_confusion_matrix
+from torchmetrics.functional.clustering.utils import check_cluster_labels
+from torchmetrics.utilities.imports import _TORCH_LINEAR_ASSIGNMENT_AVAILABLE
+
+if not _TORCH_LINEAR_ASSIGNMENT_AVAILABLE:
+    __doctest_skip__ = ["cluster_accuracy"]
+
+
+def _cluster_accuracy_compute(confmat: Tensor) -> Tensor:
+    """Computes the clustering accuracy from a confusion matrix."""
+    from torch_linear_assignment import batch_linear_assignment
+
+    confmat = confmat[None]
+    # solve the linear sum assignment problem
+    assignment = batch_linear_assignment(confmat.max() - confmat)
+    confmat = confmat[0]
+    # extract the true positives
+    tps = confmat[torch.arange(confmat.shape[0]), assignment.flatten()]
+    return tps.sum() / confmat.sum()
+
+
+def cluster_accuracy(preds: Tensor, target: Tensor, num_classes: int) -> Tensor:
+    """Computes the clustering accuracy between the predicted and target clusters.
+
+    Args:
+        preds: predicted cluster labels
+        target: ground truth cluster labels
+        num_classes: number of classes
+
+    Returns:
+        Scalar tensor with clustering accuracy between 0.0 and 1.0
+
+    Raises:
+        RuntimeError:
+            If `torch_linear_assignment` is not installed
+
+    Example:
+        >>> from torchmetrics.functional.clustering import cluster_accuracy
+        >>> preds = torch.tensor([0, 0, 1, 1])
+        >>> target = torch.tensor([1, 1, 0, 0])
+        >>> cluster_accuracy(preds, target, 2)
+        tensor(1.000)
+
+    """
+    if not _TORCH_LINEAR_ASSIGNMENT_AVAILABLE:
+        raise RuntimeError(
+            "Missing `torch_linear_assignment`. Please install it with `pip install torchmetrics[clustering]`."
+        )
+    check_cluster_labels(preds, target)
+    confmat = multiclass_confusion_matrix(preds, target, num_classes=num_classes)
+    return _cluster_accuracy_compute(confmat)

src/torchmetrics/utilities/imports.py

@@ -59,6 +59,6 @@
 _SENTENCEPIECE_AVAILABLE = RequirementCache("sentencepiece")
 _SCIPI_AVAILABLE = RequirementCache("scipy")
 _SKLEARN_GREATER_EQUAL_1_3 = RequirementCache("scikit-learn>=1.3.0")
+_TORCH_LINEAR_ASSIGNMENT_AVAILABLE = RequirementCache("torch_linear_assignment")
 _PYTDC_AVAILABLE = RequirementCache("pyTDC")
-
 _LATEX_AVAILABLE: bool = shutil.which("latex") is not None