Weighted AUROC to omit empty classes

CHANGELOG.md

@@ -46,6 +46,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Fixed
 
+- Fixed `weighted`, `multi-class` AUROC computation to allow for 0 observations of some class, as contribution to final AUROC is 0 ([#348](https://github.com/PyTorchLightning/metrics/issues/348))
 
 
 ## [0.4.1] - 2021-07-05

tests/classification/test_auroc.py

@@ -195,3 +195,31 @@ def test_error_multiclass_no_num_classes():
  ValueError, match="Detected input to `multiclass` but you did not provide `num_classes` argument"
  ):
  _ = auroc(torch.randn(20, 3).softmax(dim=-1), torch.randint(3, (20, )))
+
+
+def test_weighted_with_empty_classes():
+ """ Tests that weighted multiclass AUROC calculation yields the same results if a new
+ but empty class exists. Tests that the proper warnings and errors are raised
+ """
+ preds = torch.tensor([
+ [0.90, 0.05, 0.05],
+ [0.05, 0.90, 0.05],
+ [0.05, 0.05, 0.90],
+ [0.85, 0.05, 0.10],
+ [0.10, 0.10, 0.80],
+ ])
+ target = torch.tensor([0, 1, 1, 2, 2])
+ num_classes = 3
+ _auroc = auroc(preds, target, average="weighted", num_classes=num_classes)
+
+ # Add in a class with zero observations at second to last index
+ preds = torch.cat((preds[:, :num_classes - 1], torch.rand_like(preds[:, 0:1]), preds[:, num_classes - 1:]), axis=1)
+ # Last class (2) gets moved to 3
+ target[target == num_classes - 1] = num_classes
+ with pytest.warns(UserWarning, match='Class 2 had 0 observations, omitted from AUROC calculation'):
+ _auroc_empty_class = auroc(preds, target, average="weighted", num_classes=num_classes + 1)
+ assert _auroc == _auroc_empty_class
+
+ target = torch.zeros_like(target)
+ with pytest.raises(ValueError, match='Found 1 non-empty class in `multiclass` AUROC calculation'):
+ _ = auroc(preds, target, average="weighted", num_classes=num_classes + 1)

torchmetrics/functional/classification/auroc.py

@@ -11,6 +11,7 @@
 # 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 warnings
 from typing import Optional, Sequence, Tuple
 
 import torch
@@ -87,8 +88,23 @@ def _auroc_compute(
  else:
  raise ValueError('Detected input to be `multilabel` but you did not provide `num_classes` argument')
  else:
- if mode != DataType.BINARY and num_classes is None:
- raise ValueError('Detected input to `multiclass` but you did not provide `num_classes` argument')
+ if mode != DataType.BINARY:
+ if num_classes is None:
+ raise ValueError("Detected input to `multiclass` but you did not provide `num_classes` argument")
+ if average == AverageMethod.WEIGHTED and len(torch.unique(target)) < num_classes:
+ # If one or more classes has 0 observations, we should exclude them, as its weight will be 0
+ target_bool_mat = torch.zeros((len(target), num_classes), dtype=bool)
+ target_bool_mat[torch.arange(len(target)), target.long()] = 1
+ class_observed = target_bool_mat.sum(axis=0) > 0
+ for c in range(num_classes):
+ if not class_observed[c]:
+ warnings.warn(f'Class {c} had 0 observations, omitted from AUROC calculation', UserWarning)
+ preds = preds[:, class_observed]
+ target = target_bool_mat[:, class_observed]
+ target = torch.where(target)[1]
+ num_classes = class_observed.sum()
+ if num_classes == 1:
+ raise ValueError('Found 1 non-empty class in `multiclass` AUROC calculation')
  fpr, tpr, _ = roc(preds, target, num_classes, pos_label, sample_weights)
 
  # calculate standard roc auc score