Prevent overflow in ROC AUC and average precision calculations

monai/metrics/average_precision.py

@@ -12,7 +12,6 @@
 from __future__ import annotations
 
 import warnings
-from typing import cast
 
 import numpy as np
 import numpy.typing as npt
@@ -84,6 +83,7 @@ def aggregate(self, average: Average | str | None = None) -> np.ndarray | float
 
 
 def _calculate(y_pred: torch.Tensor, y: torch.Tensor) -> float:
+    """Compute binary average precision, grouping equal prediction scores."""
     if not (y.ndimension() == y_pred.ndimension() == 1 and len(y) == len(y_pred)):
         raise AssertionError("y and y_pred must be 1 dimension data with same length.")
     y_unique = y.unique()
@@ -103,7 +103,7 @@ def _calculate(y_pred: torch.Tensor, y: torch.Tensor) -> float:
     npos = ap = tmp_pos = 0.0
 
     for i in range(n):
-        y_i = cast(float, y[i])
+        y_i = float(y[i])
         if i + 1 < n and y_pred[i] == y_pred[i + 1]:
             tmp_pos += y_i
         else:
@@ -181,6 +181,6 @@ def compute_average_precision(
     if average == Average.MACRO:
         return np.mean(ap_values)
     if average == Average.WEIGHTED:
-        weights = [sum(y_) for y_ in y]
+        weights = np.sum(y.cpu().numpy(), axis=1, dtype=np.float64)
         return np.average(ap_values, weights=weights)  # type: ignore[no-any-return]
     raise ValueError(f'Unsupported average: {average}, available options are ["macro", "weighted", "micro", "none"].')

monai/metrics/rocauc.py

@@ -12,7 +12,6 @@
 from __future__ import annotations
 
 import warnings
-from typing import cast
 
 import numpy as np
 import numpy.typing as npt
@@ -73,6 +72,7 @@ def aggregate(self, average: Average | str | None = None) -> np.ndarray | float
 
 
 def _calculate(y_pred: torch.Tensor, y: torch.Tensor) -> float:
+    """Compute binary ROC AUC, assigning half credit to tied prediction scores."""
     if not (y.ndimension() == y_pred.ndimension() == 1 and len(y) == len(y_pred)):
         raise AssertionError("y and y_pred must be 1 dimension data with same length.")
     y_unique = y.unique()
@@ -94,7 +94,7 @@ def _calculate(y_pred: torch.Tensor, y: torch.Tensor) -> float:
     nneg = auc = tmp_pos = tmp_neg = 0.0
 
     for i in range(n):
-        y_i = cast(float, y[i])
+        y_i = float(y[i])
         if i + 1 < n and y_pred[i] == y_pred[i + 1]:
             tmp_pos += y_i
             tmp_neg += 1 - y_i
@@ -178,6 +178,6 @@ def compute_roc_auc(
     if average == Average.MACRO:
         return np.mean(auc_values)
     if average == Average.WEIGHTED:
-        weights = [sum(y_) for y_ in y]
+        weights = np.sum(y.cpu().numpy(), axis=1, dtype=np.float64)
         return np.average(auc_values, weights=weights)  # type: ignore[no-any-return]
     raise ValueError(f'Unsupported average: {average}, available options are ["macro", "weighted", "micro", "none"].')

tests/metrics/test_compute_average_precision.py

@@ -12,6 +12,7 @@
 from __future__ import annotations
 
 import unittest
+from itertools import product
 
 import numpy as np
 import torch
@@ -157,6 +158,55 @@ def test_class_value(self, y_pred, y, softmax, to_onehot, average, expected_valu
         metric.reset()
         np.testing.assert_allclose(expected_value, result, rtol=1e-5)
 
+    @parameterized.expand(
+        product(
+            ("cpu", "cuda"),
+            (torch.uint8, torch.int8, torch.int64, torch.bool, torch.float16, torch.float32, torch.float64),
+        )
+    )
+    def test_tied_scores_label_dtype(self, device, dtype):
+        """Check tied-score average precision across label dtypes on CPU and available CUDA devices."""
+        if device == "cuda" and not torch.cuda.is_available():
+            self.skipTest("CUDA is not available.")
+
+        y = torch.tensor([0] + [1] * 256, dtype=dtype, device=device)
+        y_pred = torch.tensor([1.0] + [0.0] * 256, device=device)
+        expected = 256 / 257
+
+        result = compute_average_precision(y_pred, y)
+        np.testing.assert_allclose(result, expected, rtol=1e-6)
+
+        metric = AveragePrecisionMetric()
+        for predictions, labels in zip(y_pred.split(128), y.split(128)):
+            metric(predictions, labels)
+        np.testing.assert_allclose(metric.aggregate(), expected, rtol=1e-6)
+
+    @parameterized.expand(
+        product(
+            ("cpu", "cuda"),
+            (torch.uint8, torch.int8, torch.int64, torch.bool, torch.float16, torch.float32, torch.float64),
+        )
+    )
+    def test_weighted_label_dtype(self, device, dtype):
+        """Check weighted scores with supports above float16's exact integer range."""
+        if device == "cuda" and not torch.cuda.is_available():
+            self.skipTest("CUDA is not available.")
+
+        y = torch.tensor([[1, 0]] * 3072 + [[0, 1]] * 1536, dtype=dtype, device=device)
+        y_pred = torch.zeros((4608, 2), device=device)
+        y_pred[:3072, 0] = 1
+        expected_values = [("none", [1.0, 1 / 3]), ("macro", 2 / 3), ("weighted", 7 / 9)]
+
+        for average, expected in expected_values:
+            with self.subTest(average=average):
+                result = compute_average_precision(y_pred, y, average=average)
+                np.testing.assert_allclose(result, expected, rtol=1e-6)
+
+                metric = AveragePrecisionMetric(average=average)
+                for predictions, labels in zip(y_pred.split(128), y.split(128)):
+                    metric(predictions, labels)
+                np.testing.assert_allclose(metric.aggregate(), expected, rtol=1e-6)
+
 
 if __name__ == "__main__":
     unittest.main()

tests/metrics/test_compute_roc_auc.py

@@ -12,6 +12,7 @@
 from __future__ import annotations
 
 import unittest
+from itertools import product
 
 import numpy as np
 import torch
@@ -150,6 +151,55 @@ def test_class_value(self, y_pred, y, softmax, to_onehot, average, expected_valu
         metric.reset()
         np.testing.assert_allclose(expected_value, result, rtol=1e-5)
 
+    @parameterized.expand(
+        product(
+            ("cpu", "cuda"),
+            (torch.uint8, torch.int8, torch.int64, torch.bool, torch.float16, torch.float32, torch.float64),
+        )
+    )
+    def test_tied_scores_label_dtype(self, device, dtype):
+        """Check tied-score ROC AUC across label dtypes on CPU and available CUDA devices."""
+        if device == "cuda" and not torch.cuda.is_available():
+            self.skipTest("CUDA is not available.")
+
+        y = torch.tensor([0] * 4096 + [1] * 4096, dtype=dtype, device=device)
+        y_pred = torch.zeros(y.shape, device=device)
+        expected = 0.5
+
+        result = compute_roc_auc(y_pred, y)
+        np.testing.assert_allclose(result, expected, rtol=1e-6)
+
+        metric = ROCAUCMetric()
+        for predictions, labels in zip(y_pred.split(128), y.split(128)):
+            metric(predictions, labels)
+        np.testing.assert_allclose(metric.aggregate(), expected, rtol=1e-6)
+
+    @parameterized.expand(
+        product(
+            ("cpu", "cuda"),
+            (torch.uint8, torch.int8, torch.int64, torch.bool, torch.float16, torch.float32, torch.float64),
+        )
+    )
+    def test_weighted_label_dtype(self, device, dtype):
+        """Check weighted scores with supports above float16's exact integer range."""
+        if device == "cuda" and not torch.cuda.is_available():
+            self.skipTest("CUDA is not available.")
+
+        y = torch.tensor([[1, 0]] * 3072 + [[0, 1]] * 1536, dtype=dtype, device=device)
+        y_pred = torch.zeros((4608, 2), device=device)
+        y_pred[:3072, 0] = 1
+        expected_values = [("none", [1.0, 0.5]), ("macro", 0.75), ("weighted", 5 / 6)]
+
+        for average, expected in expected_values:
+            with self.subTest(average=average):
+                result = compute_roc_auc(y_pred, y, average=average)
+                np.testing.assert_allclose(result, expected, rtol=1e-6)
+
+                metric = ROCAUCMetric(average=average)
+                for predictions, labels in zip(y_pred.split(128), y.split(128)):
+                    metric(predictions, labels)
+                np.testing.assert_allclose(metric.aggregate(), expected, rtol=1e-6)
+
 
 if __name__ == "__main__":
     unittest.main()