Convert classification wrapper to metrics (#1963)

Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Lightning-AI · Aug 3, 2023 · f6b5890 · f6b5890
1 parent ea17ffb
commit f6b5890
Show file tree

Hide file tree

Showing 67 changed files with 620 additions and 87 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -28,7 +28,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Changed
 
--
+- Changed all non-task specific classification metrics to be true subtypes of `Metric` ([#1963](https://github.com/Lightning-AI/torchmetrics/pull/1963))
 
 
 ### Removed

diff --git a/docs/source/classification/accuracy.rst b/docs/source/classification/accuracy.rst
@@ -13,7 +13,7 @@ ________________
 .. autoclass:: torchmetrics.Accuracy
     :noindex:
     :exclude-members: update, compute
-
+    :special-members: __new__
 
 BinaryAccuracy
 ^^^^^^^^^^^^^^

diff --git a/docs/source/classification/auroc.rst b/docs/source/classification/auroc.rst
@@ -15,6 +15,7 @@ ________________
 .. autoclass:: torchmetrics.AUROC
     :noindex:
     :exclude-members: update, compute
+    :special-members: __new__
 
 BinaryAUROC
 ^^^^^^^^^^^

diff --git a/docs/source/classification/average_precision.rst b/docs/source/classification/average_precision.rst
@@ -13,6 +13,7 @@ ________________
 .. autoclass:: torchmetrics.AveragePrecision
     :noindex:
     :exclude-members: update, compute
+    :special-members: __new__
 
 BinaryAveragePrecision
 ^^^^^^^^^^^^^^^^^^^^^^

diff --git a/docs/source/classification/calibration_error.rst b/docs/source/classification/calibration_error.rst
@@ -15,6 +15,7 @@ ________________
 .. autoclass:: torchmetrics.CalibrationError
     :noindex:
     :exclude-members: update, compute
+    :special-members: __new__
 
 BinaryCalibrationError
 ^^^^^^^^^^^^^^^^^^^^^^

diff --git a/docs/source/classification/cohen_kappa.rst b/docs/source/classification/cohen_kappa.rst
@@ -12,12 +12,10 @@ Cohen Kappa
 Module Interface
 ________________
 
-CohenKappa
-^^^^^^^^^^
-
 .. autoclass:: torchmetrics.CohenKappa
     :noindex:
     :exclude-members: update, compute
+    :special-members: __new__
 
 BinaryCohenKappa
 ^^^^^^^^^^^^^^^^

diff --git a/docs/source/classification/confusion_matrix.rst b/docs/source/classification/confusion_matrix.rst
@@ -12,12 +12,10 @@ Confusion Matrix
 Module Interface
 ________________
 
-ConfusionMatrix
-^^^^^^^^^^^^^^^
-
 .. autoclass:: torchmetrics.ConfusionMatrix
     :noindex:
     :exclude-members: update, compute
+    :special-members: __new__
 
 BinaryConfusionMatrix
 ^^^^^^^^^^^^^^^^^^^^^

diff --git a/docs/source/classification/exact_match.rst b/docs/source/classification/exact_match.rst
@@ -10,12 +10,10 @@ Exact Match
 Module Interface
 ________________
 
-ExactMatch
-^^^^^^^^^^^^^^^
-
 .. autoclass:: torchmetrics.ExactMatch
     :noindex:
     :exclude-members: update, compute
+    :special-members: __new__
 
 MulticlassExactMatch
 ^^^^^^^^^^^^^^^^^^^^

diff --git a/docs/source/classification/f1_score.rst b/docs/source/classification/f1_score.rst
@@ -10,12 +10,10 @@ F-1 Score
 Module Interface
 ________________
 
-F1Score
-^^^^^^^
-
 .. autoclass:: torchmetrics.F1Score
     :noindex:
     :exclude-members: update, compute
+    :special-members: __new__
 
 BinaryF1Score
 ^^^^^^^^^^^^^

diff --git a/docs/source/classification/fbeta_score.rst b/docs/source/classification/fbeta_score.rst
@@ -12,12 +12,10 @@ F-Beta Score
 Module Interface
 ________________
 
-FBetaScore
-^^^^^^^^^^
-
 .. autoclass:: torchmetrics.FBetaScore
     :noindex:
     :exclude-members: update, compute
+    :special-members: __new__
 
 BinaryFBetaScore
 ^^^^^^^^^^^^^^^^

diff --git a/docs/source/classification/hamming_distance.rst b/docs/source/classification/hamming_distance.rst
@@ -10,12 +10,10 @@ Hamming Distance
 Module Interface
 ________________
 
-HammingDistance
-^^^^^^^^^^^^^^^
-
 .. autoclass:: torchmetrics.HammingDistance
     :noindex:
     :exclude-members: update, compute
+    :special-members: __new__
 
 BinaryHammingDistance
 ^^^^^^^^^^^^^^^^^^^^^

diff --git a/docs/source/classification/hinge_loss.rst b/docs/source/classification/hinge_loss.rst
@@ -13,6 +13,7 @@ ________________
 .. autoclass:: torchmetrics.HingeLoss
     :noindex:
     :exclude-members: update, compute
+    :special-members: __new__
 
 BinaryHingeLoss
 ^^^^^^^^^^^^^^^

diff --git a/docs/source/classification/jaccard_index.rst b/docs/source/classification/jaccard_index.rst
@@ -10,12 +10,10 @@ Jaccard Index
 Module Interface
 ________________
 
-JaccardIndex
-^^^^^^^^^^^^
-
 .. autoclass:: torchmetrics.JaccardIndex
     :noindex:
     :exclude-members: update, compute
+    :special-members: __new__
 
 BinaryJaccardIndex
 ^^^^^^^^^^^^^^^^^^

diff --git a/docs/source/classification/matthews_corr_coef.rst b/docs/source/classification/matthews_corr_coef.rst
@@ -12,12 +12,10 @@ Matthews Correlation Coefficient
 Module Interface
 ________________
 
-MatthewsCorrCoef
-^^^^^^^^^^^^^^^^
-
 .. autoclass:: torchmetrics.MatthewsCorrCoef
     :noindex:
     :exclude-members: update, compute
+    :special-members: __new__
 
 BinaryMatthewsCorrCoef
 ^^^^^^^^^^^^^^^^^^^^^^

diff --git a/docs/source/classification/precision.rst b/docs/source/classification/precision.rst
@@ -15,6 +15,7 @@ ________________
 .. autoclass:: torchmetrics.Precision
     :noindex:
     :exclude-members: update, compute
+    :special-members: __new__
 
 BinaryPrecision
 ^^^^^^^^^^^^^^^

diff --git a/docs/source/classification/precision_at_fixed_recall.rst b/docs/source/classification/precision_at_fixed_recall.rst
@@ -10,6 +10,11 @@ Precision At Fixed Recall
 Module Interface
 ________________
 
+.. autoclass:: torchmetrics.PrecisionAtFixedRecall
+    :noindex:
+    :exclude-members: update, compute
+    :special-members: __new__
+
 BinaryPrecisionAtFixedRecall
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 

diff --git a/docs/source/classification/precision_recall_curve.rst b/docs/source/classification/precision_recall_curve.rst
@@ -13,6 +13,7 @@ ________________
 .. autoclass:: torchmetrics.PrecisionRecallCurve
     :noindex:
     :exclude-members: update, compute
+    :special-members: __new__
 
 BinaryPrecisionRecallCurve
 ^^^^^^^^^^^^^^^^^^^^^^^^^^

diff --git a/docs/source/classification/recall.rst b/docs/source/classification/recall.rst
@@ -13,6 +13,7 @@ ________________
 .. autoclass:: torchmetrics.Recall
     :noindex:
     :exclude-members: update, compute
+    :special-members: __new__
 
 BinaryRecall
 ^^^^^^^^^^^^

diff --git a/docs/source/classification/recall_at_fixed_precision.rst b/docs/source/classification/recall_at_fixed_precision.rst
@@ -10,6 +10,11 @@ Recall At Fixed Precision
 Module Interface
 ________________
 
+.. autoclass:: torchmetrics.RecallAtFixedPrecision
+    :noindex:
+    :exclude-members: update, compute
+    :special-members: __new__
+
 BinaryRecallAtFixedPrecision
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 

diff --git a/docs/source/classification/roc.rst b/docs/source/classification/roc.rst
@@ -13,6 +13,7 @@ ________________
 .. autoclass:: torchmetrics.ROC
     :noindex:
     :exclude-members: update, compute
+    :special-members: __new__
 
 BinaryROC
 ^^^^^^^^^

diff --git a/docs/source/classification/specificity.rst b/docs/source/classification/specificity.rst
@@ -13,6 +13,7 @@ ________________
 .. autoclass:: torchmetrics.Specificity
     :noindex:
     :exclude-members: update, compute
+    :special-members: __new__
 
 BinarySpecificity
 ^^^^^^^^^^^^^^^^^

diff --git a/docs/source/classification/specificity_at_sensitivity.rst b/docs/source/classification/specificity_at_sensitivity.rst
@@ -10,6 +10,11 @@ Specificity At Sensitivity
 Module Interface
 ________________
 
+.. autoclass:: torchmetrics.SpecificityAtSensitivity
+    :noindex:
+    :exclude-members: update, compute
+    :special-members: __new__
+
 BinarySpecificityAtSensitivity
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 

diff --git a/docs/source/classification/stat_scores.rst b/docs/source/classification/stat_scores.rst
@@ -12,12 +12,10 @@ Stat Scores
 Module Interface
 ________________
 
-StatScores
-^^^^^^^^^^
-
 .. autoclass:: torchmetrics.StatScores
     :noindex:
     :exclude-members: update, compute
+    :special-members: __new__
 
 BinaryStatScores
 ^^^^^^^^^^^^^^^^

diff --git a/docs/source/conf.py b/docs/source/conf.py
@@ -17,11 +17,6 @@
 import shutil
 import sys
 
-import torch
-
-# this removes "Initializes internal Module state, shared by both nn.Module and ScriptModule." from the docs
-torch.nn.Module.__init__.__doc__ = ""
-
 import pt_lightning_sphinx_theme
 from lightning_utilities.docs import fetch_external_assets
 from lightning_utilities.docs.formatting import _transform_changelog
@@ -391,7 +386,7 @@ def _get_version_str():
 
 autodoc_member_order = "groupwise"
 
-autoclass_content = "both"
+autoclass_content = "class"
 
 autodoc_default_options = {
     "members": True,

diff --git a/src/torchmetrics/__init__.py b/src/torchmetrics/__init__.py
@@ -52,6 +52,7 @@
     Recall,
     RecallAtFixedPrecision,
     Specificity,
+    SpecificityAtSensitivity,
     StatScores,
 )
 from torchmetrics.collections import MetricCollection  # noqa: E402
@@ -223,6 +224,7 @@
     "SignalNoiseRatio",
     "SpearmanCorrCoef",
     "Specificity",
+    "SpecificityAtSensitivity",
     "SpectralAngleMapper",
     "SpectralDistortionIndex",
     "SQuAD",

diff --git a/src/torchmetrics/classification/__init__.py b/src/torchmetrics/classification/__init__.py
@@ -107,6 +107,7 @@
     BinarySpecificityAtSensitivity,
     MulticlassSpecificityAtSensitivity,
     MultilabelSpecificityAtSensitivity,
+    SpecificityAtSensitivity,
 )
 from torchmetrics.classification.stat_scores import (
     BinaryStatScores,
@@ -201,6 +202,7 @@
     "MulticlassSpecificityAtSensitivity",
     "MultilabelSpecificityAtSensitivity",
     "BinaryPrecisionAtFixedRecall",
+    "SpecificityAtSensitivity",
     "MulticlassPrecisionAtFixedRecall",
     "MultilabelPrecisionAtFixedRecall",
     "PrecisionAtFixedRecall",

diff --git a/src/torchmetrics/classification/accuracy.py b/src/torchmetrics/classification/accuracy.py
@@ -16,6 +16,7 @@
 from torch import Tensor
 from typing_extensions import Literal
 
+from torchmetrics.classification.base import _ClassificationTaskWrapper
 from torchmetrics.classification.stat_scores import BinaryStatScores, MulticlassStatScores, MultilabelStatScores
 from torchmetrics.functional.classification.accuracy import _accuracy_reduce
 from torchmetrics.metric import Metric
@@ -445,7 +446,7 @@ def plot(
         return self._plot(val, ax)
 
 
-class Accuracy:
+class Accuracy(_ClassificationTaskWrapper):
     r"""Compute `Accuracy`_.
 
     .. math::

diff --git a/src/torchmetrics/classification/auroc.py b/src/torchmetrics/classification/auroc.py
@@ -16,6 +16,7 @@
 from torch import Tensor
 from typing_extensions import Literal
 
+from torchmetrics.classification.base import _ClassificationTaskWrapper
 from torchmetrics.classification.precision_recall_curve import (
     BinaryPrecisionRecallCurve,
     MulticlassPrecisionRecallCurve,
@@ -460,7 +461,7 @@ def plot(  # type: ignore[override]
         return self._plot(val, ax)
 
 
-class AUROC:
+class AUROC(_ClassificationTaskWrapper):
     r"""Compute Area Under the Receiver Operating Characteristic Curve (`ROC AUC`_).
 
     The AUROC score summarizes the ROC curve into an single number that describes the performance of a model for
@@ -518,3 +519,15 @@ def __new__(  # type: ignore[misc]
                 raise ValueError(f"`num_labels` is expected to be `int` but `{type(num_labels)} was passed.`")
             return MultilabelAUROC(num_labels, average, **kwargs)
         raise ValueError(f"Task {task} not supported!")
+
+    def update(self, *args: Any, **kwargs: Any) -> None:
+        """Update metric state."""
+        raise NotImplementedError(
+            f"{self.__class__.__name__} metric does not have a global `update` method. Use the task specific metric."
+        )
+
+    def compute(self) -> None:
+        """Compute metric."""
+        raise NotImplementedError(
+            f"{self.__class__.__name__} metric does not have a global `compute` method. Use the task specific metric."
+        )
diff --git a/src/torchmetrics/classification/average_precision.py b/src/torchmetrics/classification/average_precision.py
@@ -16,6 +16,7 @@
 from torch import Tensor
 from typing_extensions import Literal
 
+from torchmetrics.classification.base import _ClassificationTaskWrapper
 from torchmetrics.classification.precision_recall_curve import (
     BinaryPrecisionRecallCurve,
     MulticlassPrecisionRecallCurve,
@@ -465,7 +466,7 @@ def plot(  # type: ignore[override]
         return self._plot(val, ax)
 
 
-class AveragePrecision:
+class AveragePrecision(_ClassificationTaskWrapper):
     r"""Compute the average precision (AP) score.
 
     The AP score summarizes a precision-recall curve as an weighted mean of precisions at each threshold, with the