Fix missing AUC score when plotting (#1948)

(cherry picked from commit 7cda67c)
Lightning-AI · Aug 3, 2023 · 256a6d6 · 256a6d6
1 parent 903071d
commit 256a6d6
Show file tree

Hide file tree

Showing 3 changed files with 85 additions and 60 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -26,6 +26,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Fixed bug related to the `prefix/postfix` arguments in `MetricCollection` and `ClasswiseWrapper` being duplicated ([#1918](https://github.com/Lightning-AI/torchmetrics/pull/1918))
 
 
+- Fixed missing AUC score when plotting classification metrics that support the `score` argument ([#1948](https://github.com/Lightning-AI/torchmetrics/pull/1948))
+
+
 ## [1.0.1] - 2023-07-13
 
 ### Fixed

diff --git a/src/torchmetrics/classification/precision_recall_curve.py b/src/torchmetrics/classification/precision_recall_curve.py
@@ -196,18 +196,24 @@ def plot(
         .. plot::
             :scale: 75
 
-            >>> from torch import randn, randint
-            >>> import torch.nn.functional as F
-            >>> from torchmetrics.classification import BinaryROC
-            >>> preds = F.softmax(randn(20, 2), dim=1)
+            >>> from torch import rand, randint
+            >>> from torchmetrics.classification import BinaryPrecisionRecallCurve
+            >>> preds = rand(20)
             >>> target = randint(2, (20,))
-            >>> metric = BinaryROC()
-            >>> metric.update(preds[:, 1], target)
-            >>> fig_, ax_ = metric.plot()
+            >>> metric = BinaryPrecisionRecallCurve()
+            >>> metric.update(preds, target)
+            >>> fig_, ax_ = metric.plot(score=True)
+
         """
-        curve = curve or self.compute()
-        score = _auc_compute_without_check(curve[0], curve[1], 1.0) if not curve and score is True else None
-        return plot_curve(curve, score=score, ax=ax, label_names=("Precision", "Recall"), name=self.__class__.__name__)
+        curve_computed = curve or self.compute()
+        score = (
+            _auc_compute_without_check(curve_computed[0], curve_computed[1], 1.0)
+            if not curve and score is True
+            else None
+        )
+        return plot_curve(
+            curve_computed, score=score, ax=ax, label_names=("Precision", "Recall"), name=self.__class__.__name__
+        )
 
 
 class MulticlassPrecisionRecallCurve(Metric):
@@ -371,17 +377,21 @@ def plot(
             :scale: 75
 
             >>> from torch import randn, randint
-            >>> import torch.nn.functional as F
-            >>> from torchmetrics.classification import BinaryROC
-            >>> preds = F.softmax(randn(20, 2), dim=1)
-            >>> target = randint(2, (20,))
-            >>> metric = BinaryROC()
-            >>> metric.update(preds[:, 1], target)
-            >>> fig_, ax_ = metric.plot()
+            >>> from torchmetrics.classification import MulticlassPrecisionRecallCurve
+            >>> preds = randn(20, 3).softmax(dim=-1)
+            >>> target = randint(3, (20,))
+            >>> metric = MulticlassPrecisionRecallCurve(num_classes=3)
+            >>> metric.update(preds, target)
+            >>> fig_, ax_ = metric.plot(score=True)
+
         """
-        curve = curve or self.compute()
-        score = _reduce_auroc(curve[0], curve[1], average=None) if not curve and score is True else None
-        return plot_curve(curve, score=score, ax=ax, label_names=("Precision", "Recall"), name=self.__class__.__name__)
+        curve_computed = curve or self.compute()
+        score = (
+            _reduce_auroc(curve_computed[0], curve_computed[1], average=None) if not curve and score is True else None
+        )
+        return plot_curve(
+            curve_computed, score=score, ax=ax, label_names=("Precision", "Recall"), name=self.__class__.__name__
+        )
 
 
 class MultilabelPrecisionRecallCurve(Metric):
@@ -553,18 +563,22 @@ def plot(
         .. plot::
             :scale: 75
 
-            >>> from torch import randn, randint
-            >>> import torch.nn.functional as F
-            >>> from torchmetrics.classification import BinaryROC
-            >>> preds = F.softmax(randn(20, 2), dim=1)
-            >>> target = randint(2, (20,))
-            >>> metric = BinaryROC()
-            >>> metric.update(preds[:, 1], target)
-            >>> fig_, ax_ = metric.plot()
+            >>> from torch import rand, randint
+            >>> from torchmetrics.classification import MultilabelPrecisionRecallCurve
+            >>> preds = rand(20, 3)
+            >>> target = randint(2, (20,3))
+            >>> metric = MultilabelPrecisionRecallCurve(num_labels=3)
+            >>> metric.update(preds, target)
+            >>> fig_, ax_ = metric.plot(score=True)
+
         """
-        curve = curve or self.compute()
-        score = _reduce_auroc(curve[0], curve[1], average=None) if not curve and score is True else None
-        return plot_curve(curve, score=score, ax=ax, label_names=("Precision", "Recall"), name=self.__class__.__name__)
+        curve_computed = curve or self.compute()
+        score = (
+            _reduce_auroc(curve_computed[0], curve_computed[1], average=None) if not curve and score is True else None
+        )
+        return plot_curve(
+            curve_computed, score=score, ax=ax, label_names=("Precision", "Recall"), name=self.__class__.__name__
+        )
 
 
 class PrecisionRecallCurve:

diff --git a/src/torchmetrics/classification/roc.py b/src/torchmetrics/classification/roc.py
@@ -142,19 +142,23 @@ def plot(
         .. plot::
             :scale: 75
 
-            >>> from torch import randn, randint
-            >>> import torch.nn.functional as F
+            >>> from torch import rand, randint
             >>> from torchmetrics.classification import BinaryROC
-            >>> preds = F.softmax(randn(20, 2), dim=1)
+            >>> preds = rand(20)
             >>> target = randint(2, (20,))
             >>> metric = BinaryROC()
-            >>> metric.update(preds[:, 1], target)
-            >>> fig_, ax_ = metric.plot()
+            >>> metric.update(preds, target)
+            >>> fig_, ax_ = metric.plot(score=True)
+
         """
-        curve = curve or self.compute()
-        score = _auc_compute_without_check(curve[0], curve[1], 1.0) if not curve and score is True else None
+        curve_computed = curve or self.compute()
+        score = (
+            _auc_compute_without_check(curve_computed[0], curve_computed[1], 1.0)
+            if not curve and score is True
+            else None
+        )
         return plot_curve(
-            curve,
+            curve_computed,
             score=score,
             ax=ax,
             label_names=("False positive rate", "True positive rate"),
@@ -293,18 +297,20 @@ def plot(
             :scale: 75
 
             >>> from torch import randn, randint
-            >>> import torch.nn.functional as F
-            >>> from torchmetrics.classification import BinaryROC
-            >>> preds = F.softmax(randn(20, 2), dim=1)
-            >>> target = randint(2, (20,))
-            >>> metric = BinaryROC()
-            >>> metric.update(preds[:, 1], target)
-            >>> fig_, ax_ = metric.plot()
+            >>> from torchmetrics.classification import MulticlassROC
+            >>> preds = randn(20, 3).softmax(dim=-1)
+            >>> target = randint(3, (20,))
+            >>> metric = MulticlassROC(num_classes=3)
+            >>> metric.update(preds, target)
+            >>> fig_, ax_ = metric.plot(score=True)
+
         """
-        curve = curve or self.compute()
-        score = _reduce_auroc(curve[0], curve[1], average=None) if not curve and score is True else None
+        curve_computed = curve or self.compute()
+        score = (
+            _reduce_auroc(curve_computed[0], curve_computed[1], average=None) if not curve and score is True else None
+        )
         return plot_curve(
-            curve,
+            curve_computed,
             score=score,
             ax=ax,
             label_names=("False positive rate", "True positive rate"),
@@ -444,19 +450,21 @@ def plot(
         .. plot::
             :scale: 75
 
-            >>> from torch import randn, randint
-            >>> import torch.nn.functional as F
-            >>> from torchmetrics.classification import BinaryROC
-            >>> preds = F.softmax(randn(20, 2), dim=1)
-            >>> target = randint(2, (20,))
-            >>> metric = BinaryROC()
-            >>> metric.update(preds[:, 1], target)
-            >>> fig_, ax_ = metric.plot()
+            >>> from torch import rand, randint
+            >>> from torchmetrics.classification import MultilabelROC
+            >>> preds = rand(20, 3)
+            >>> target = randint(2, (20,3))
+            >>> metric = MultilabelROC(num_labels=3)
+            >>> metric.update(preds, target)
+            >>> fig_, ax_ = metric.plot(score=True)
+
         """
-        curve = curve or self.compute()
-        score = _reduce_auroc(curve[0], curve[1], average=None) if not curve and score is True else None
+        curve_computed = curve or self.compute()
+        score = (
+            _reduce_auroc(curve_computed[0], curve_computed[1], average=None) if not curve and score is True else None
+        )
         return plot_curve(
-            curve,
+            curve_computed,
             score=score,
             ax=ax,
             label_names=("False positive rate", "True positive rate"),