diff --git a/monai/metrics/rocauc.py b/monai/metrics/rocauc.py index c36bb507c9..73f15534a9 100644 --- a/monai/metrics/rocauc.py +++ b/monai/metrics/rocauc.py @@ -66,9 +66,14 @@ def aggregate(self): # type: ignore def _calculate(y_pred: torch.Tensor, y: torch.Tensor) -> float: 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.") - if not y.unique().equal(torch.tensor([0, 1], dtype=y.dtype, device=y.device)): - warnings.warn("y values must be 0 or 1, can not be all 0 or all 1, skip AUC computation and return `Nan`.") + y_unique = y.unique() + if len(y_unique) == 1: + warnings.warn(f"y values can not be all {y_unique.item()}, skip AUC computation and return `Nan`.") return float("nan") + if not y_unique.equal(torch.tensor([0, 1], dtype=y.dtype, device=y.device)): + warnings.warn(f"y values must be 0 or 1, but in {y_unique.tolist()}, skip AUC computation and return `Nan`.") + return float("nan") + n = len(y) indices = y_pred.argsort() y = y[indices].cpu().numpy() diff --git a/tests/test_compute_roc_auc.py b/tests/test_compute_roc_auc.py index 592bf9459c..887db08c7c 100644 --- a/tests/test_compute_roc_auc.py +++ b/tests/test_compute_roc_auc.py @@ -77,10 +77,39 @@ float("nan"), ] +TEST_CASE_9 = [ + torch.tensor([[0.1, 0.9], [0.3, 1.4], [0.2, 0.1], [0.1, 0.5]]), + torch.tensor([[1], [1], [1], [1]]), + True, + 2, + "macro", + float("nan"), +] + +TEST_CASE_10 = [ + torch.tensor([[0.1, 0.9], [0.3, 1.4], [0.2, 0.1], [0.1, 0.5]]), + torch.tensor([[0, 0], [1, 1], [2, 2], [3, 3]]), + True, + None, + "macro", + float("nan"), +] + class TestComputeROCAUC(unittest.TestCase): @parameterized.expand( - [TEST_CASE_1, TEST_CASE_2, TEST_CASE_3, TEST_CASE_4, TEST_CASE_5, TEST_CASE_6, TEST_CASE_7, TEST_CASE_8] + [ + TEST_CASE_1, + TEST_CASE_2, + TEST_CASE_3, + TEST_CASE_4, + TEST_CASE_5, + TEST_CASE_6, + TEST_CASE_7, + TEST_CASE_8, + TEST_CASE_9, + TEST_CASE_10, + ] ) def test_value(self, y_pred, y, softmax, to_onehot, average, expected_value): y_pred_trans = Compose([ToTensor(), Activations(softmax=softmax)]) @@ -91,7 +120,18 @@ def test_value(self, y_pred, y, softmax, to_onehot, average, expected_value): np.testing.assert_allclose(expected_value, result, rtol=1e-5) @parameterized.expand( - [TEST_CASE_1, TEST_CASE_2, TEST_CASE_3, TEST_CASE_4, TEST_CASE_5, TEST_CASE_6, TEST_CASE_7, TEST_CASE_8] + [ + TEST_CASE_1, + TEST_CASE_2, + TEST_CASE_3, + TEST_CASE_4, + TEST_CASE_5, + TEST_CASE_6, + TEST_CASE_7, + TEST_CASE_8, + TEST_CASE_9, + TEST_CASE_10, + ] ) def test_class_value(self, y_pred, y, softmax, to_onehot, average, expected_value): y_pred_trans = Compose([ToTensor(), Activations(softmax=softmax)])