fix ConfusionMatrix and StatScores for num_classes > 16 (#1521)

* fix: ConfusionMatrix&StatScores for num_classes > 16

e.g. if preds or target is uint8 and num_classes > 16, unique_mapping
overflows

* unittest #1521, NUM_CLASSES=17, add multiclass case "single dim int8-logits"

* revert tests

* add byte testing

* changelog

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Nicki Skafte Detlefsen <[email protected]>
Co-authored-by: Jirka Borovec <[email protected]>

(cherry picked from commit 6bc249d)

Author: vincentvaroquauxads

CHANGELOG.md

@@ -23,6 +23,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Fixed `multilabel` in `ExactMatch` ([#1474](https://github.com/Lightning-AI/metrics/pull/1474))
 
 
+- Fixed classification metrics for `byte` input ([#1521](https://github.com/Lightning-AI/metrics/pull/1474))
+
+
 ## [0.11.1] - 2023-01-30
 
 ### Fixed

src/torchmetrics/functional/classification/confusion_matrix.py

@@ -325,8 +325,8 @@ def _multiclass_confusion_matrix_format(
 
 
 def _multiclass_confusion_matrix_update(preds: Tensor, target: Tensor, num_classes: int) -> Tensor:
-    """Computes the bins to update the confusion matrix with."""
-    unique_mapping = (target * num_classes + preds).to(torch.long)
+    """Compute the bins to update the confusion matrix with."""
+    unique_mapping = target.to(torch.long) * num_classes + preds.to(torch.long)
     bins = _bincount(unique_mapping, minlength=num_classes**2)
     return bins.reshape(num_classes, num_classes)
 

src/torchmetrics/functional/classification/stat_scores.py

@@ -408,7 +408,7 @@ def _multiclass_stat_scores_update(
             idx = target != ignore_index
             preds = preds[idx]
             target = target[idx]
-        unique_mapping = (target * num_classes + preds).to(torch.long)
+        unique_mapping = target.to(torch.long) * num_classes + preds.to(torch.long)
         bins = _bincount(unique_mapping, minlength=num_classes**2)
         confmat = bins.reshape(num_classes, num_classes)
         tp = confmat.diag()

tests/unittests/classification/test_confusion_matrix.py

@@ -214,6 +214,18 @@ def test_multiclass_confusion_matrix_dtype_gpu(self, input, dtype):
         )
 
 
+def test_multiclass_overflow():
+    """Test that multiclass computations does not overflow even on byte input."""
+    preds = torch.randint(20, (100,)).byte()
+    target = torch.randint(20, (100,)).byte()
+
+    m = MulticlassConfusionMatrix(num_classes=20)
+    res = m(preds, target)
+
+    compare = sk_confusion_matrix(target, preds)
+    assert torch.allclose(res, torch.tensor(compare))
+
+
 def _sk_confusion_matrix_multilabel(preds, target, normalize=None, ignore_index=None):
     preds = preds.numpy()
     target = target.numpy()

tests/unittests/classification/test_stat_scores.py

@@ -323,6 +323,24 @@ def test_top_k_multiclass(k, preds, target, average, expected):
     )
 
 
+def test_multiclass_overflow():
+    """Test that multiclass computations does not overflow even on byte input."""
+    preds = torch.randint(20, (100,)).byte()
+    target = torch.randint(20, (100,)).byte()
+
+    m = MulticlassStatScores(num_classes=20, average=None)
+    res = m(preds, target)
+
+    confmat = sk_confusion_matrix(target, preds)
+    fp = confmat.sum(axis=0) - np.diag(confmat)
+    fn = confmat.sum(axis=1) - np.diag(confmat)
+    tp = np.diag(confmat)
+    tn = confmat.sum() - (fp + fn + tp)
+    compare = np.stack([tp, fp, tn, fn, tp + fn]).T
+
+    assert torch.allclose(res, torch.tensor(compare))
+
+
 def _sk_stat_scores_multilabel(preds, target, ignore_index, multidim_average, average):
     preds = preds.numpy()
     target = target.numpy()