Reduce memory in classification metrics when average='micro'

CHANGELOG.md

@@ -44,6 +44,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Fixed type casting in `MAP` metric between `bool` and `float32` ([#1150](https://github.com/Lightning-AI/metrics/pull/1150))
 
 
+- Fixed high memory usage for certain classification metrics when `average='micro'` ([#1286](https://github.com/Lightning-AI/metrics/pull/1286))
+
+
 ## [0.10.0] - 2022-10-04
 
 ### Added

src/torchmetrics/classification/stat_scores.py

@@ -44,7 +44,11 @@
 
 class _AbstractStatScores(Metric):
     # define common functions
-    def _create_state(self, size: int, multidim_average: str) -> None:
+    def _create_state(
+        self,
+        size: int,
+        multidim_average: Literal["global", "samplewise"] = "global",
+    ) -> None:
         """Initialize the states for the different statistics."""
         default: Union[Callable[[], list], Callable[[], Tensor]]
         if multidim_average == "samplewise":
@@ -53,6 +57,7 @@ def _create_state(self, size: int, multidim_average: str) -> None:
         else:
             default = lambda: torch.zeros(size, dtype=torch.long)
             dist_reduce_fx = "sum"
+
         self.add_state("tp", default(), dist_reduce_fx=dist_reduce_fx)
         self.add_state("fp", default(), dist_reduce_fx=dist_reduce_fx)
         self.add_state("tn", default(), dist_reduce_fx=dist_reduce_fx)
@@ -159,7 +164,7 @@ def __init__(
         self.ignore_index = ignore_index
         self.validate_args = validate_args
 
-        self._create_state(1, multidim_average)
+        self._create_state(size=1, multidim_average=multidim_average)
 
     def update(self, preds: Tensor, target: Tensor) -> None:  # type: ignore
         """Update state with predictions and targets.
@@ -300,7 +305,9 @@ def __init__(
         self.ignore_index = ignore_index
         self.validate_args = validate_args
 
-        self._create_state(num_classes, multidim_average)
+        self._create_state(
+            size=1 if (average == "micro" and top_k == 1) else num_classes, multidim_average=multidim_average
+        )
 
     def update(self, preds: Tensor, target: Tensor) -> None:  # type: ignore
         """Update state with predictions and targets.
@@ -315,7 +322,7 @@ def update(self, preds: Tensor, target: Tensor) -> None:  # type: ignore
             )
         preds, target = _multiclass_stat_scores_format(preds, target, self.top_k)
         tp, fp, tn, fn = _multiclass_stat_scores_update(
-            preds, target, self.num_classes, self.top_k, self.multidim_average, self.ignore_index
+            preds, target, self.num_classes, self.top_k, self.average, self.multidim_average, self.ignore_index
         )
         self._update_state(tp, fp, tn, fn)
 
@@ -448,7 +455,7 @@ def __init__(
         self.ignore_index = ignore_index
         self.validate_args = validate_args
 
-        self._create_state(num_labels, multidim_average)
+        self._create_state(size=num_labels, multidim_average=multidim_average)
 
     def update(self, preds: Tensor, target: Tensor) -> None:  # type: ignore
         """Update state with predictions and targets.

src/torchmetrics/functional/classification/accuracy.py

@@ -273,7 +273,9 @@ def multiclass_accuracy(
         _multiclass_stat_scores_arg_validation(num_classes, top_k, average, multidim_average, ignore_index)
         _multiclass_stat_scores_tensor_validation(preds, target, num_classes, multidim_average, ignore_index)
     preds, target = _multiclass_stat_scores_format(preds, target, top_k)
-    tp, fp, tn, fn = _multiclass_stat_scores_update(preds, target, num_classes, top_k, multidim_average, ignore_index)
+    tp, fp, tn, fn = _multiclass_stat_scores_update(
+        preds, target, num_classes, top_k, average, multidim_average, ignore_index
+    )
     return _accuracy_reduce(tp, fp, tn, fn, average=average, multidim_average=multidim_average)
 
 

src/torchmetrics/functional/classification/f_beta.py

@@ -275,7 +275,9 @@ def multiclass_fbeta_score(
         _multiclass_fbeta_score_arg_validation(beta, num_classes, top_k, average, multidim_average, ignore_index)
         _multiclass_stat_scores_tensor_validation(preds, target, num_classes, multidim_average, ignore_index)
     preds, target = _multiclass_stat_scores_format(preds, target, top_k)
-    tp, fp, tn, fn = _multiclass_stat_scores_update(preds, target, num_classes, top_k, multidim_average, ignore_index)
+    tp, fp, tn, fn = _multiclass_stat_scores_update(
+        preds, target, num_classes, top_k, average, multidim_average, ignore_index
+    )
     return _fbeta_reduce(tp, fp, tn, fn, beta, average=average, multidim_average=multidim_average)
 
 

src/torchmetrics/functional/classification/hamming.py

@@ -274,7 +274,9 @@ def multiclass_hamming_distance(
         _multiclass_stat_scores_arg_validation(num_classes, top_k, average, multidim_average, ignore_index)
         _multiclass_stat_scores_tensor_validation(preds, target, num_classes, multidim_average, ignore_index)
     preds, target = _multiclass_stat_scores_format(preds, target, top_k)
-    tp, fp, tn, fn = _multiclass_stat_scores_update(preds, target, num_classes, top_k, multidim_average, ignore_index)
+    tp, fp, tn, fn = _multiclass_stat_scores_update(
+        preds, target, num_classes, top_k, average, multidim_average, ignore_index
+    )
     return _hamming_distance_reduce(tp, fp, tn, fn, average=average, multidim_average=multidim_average)
 
 

src/torchmetrics/functional/classification/precision_recall.py

@@ -249,7 +249,9 @@ def multiclass_precision(
         _multiclass_stat_scores_arg_validation(num_classes, top_k, average, multidim_average, ignore_index)
         _multiclass_stat_scores_tensor_validation(preds, target, num_classes, multidim_average, ignore_index)
     preds, target = _multiclass_stat_scores_format(preds, target, top_k)
-    tp, fp, tn, fn = _multiclass_stat_scores_update(preds, target, num_classes, top_k, multidim_average, ignore_index)
+    tp, fp, tn, fn = _multiclass_stat_scores_update(
+        preds, target, num_classes, top_k, average, multidim_average, ignore_index
+    )
     return _precision_recall_reduce("precision", tp, fp, tn, fn, average=average, multidim_average=multidim_average)
 
 
@@ -542,7 +544,9 @@ def multiclass_recall(
         _multiclass_stat_scores_arg_validation(num_classes, top_k, average, multidim_average, ignore_index)
         _multiclass_stat_scores_tensor_validation(preds, target, num_classes, multidim_average, ignore_index)
     preds, target = _multiclass_stat_scores_format(preds, target, top_k)
-    tp, fp, tn, fn = _multiclass_stat_scores_update(preds, target, num_classes, top_k, multidim_average, ignore_index)
+    tp, fp, tn, fn = _multiclass_stat_scores_update(
+        preds, target, num_classes, top_k, average, multidim_average, ignore_index
+    )
     return _precision_recall_reduce("recall", tp, fp, tn, fn, average=average, multidim_average=multidim_average)
 
 

src/torchmetrics/functional/classification/specificity.py

@@ -246,7 +246,9 @@ def multiclass_specificity(
         _multiclass_stat_scores_arg_validation(num_classes, top_k, average, multidim_average, ignore_index)
         _multiclass_stat_scores_tensor_validation(preds, target, num_classes, multidim_average, ignore_index)
     preds, target = _multiclass_stat_scores_format(preds, target, top_k)
-    tp, fp, tn, fn = _multiclass_stat_scores_update(preds, target, num_classes, top_k, multidim_average, ignore_index)
+    tp, fp, tn, fn = _multiclass_stat_scores_update(
+        preds, target, num_classes, top_k, average, multidim_average, ignore_index
+    )
     return _specificity_reduce(tp, fp, tn, fn, average=average, multidim_average=multidim_average)
 
 

src/torchmetrics/functional/classification/stat_scores.py

@@ -351,6 +351,7 @@ def _multiclass_stat_scores_update(
     target: Tensor,
     num_classes: int,
     top_k: int = 1,
+    average: Optional[Literal["micro", "macro", "weighted", "none"]] = "macro",
     multidim_average: Literal["global", "samplewise"] = "global",
     ignore_index: Optional[int] = None,
 ) -> Tuple[Tensor, Tensor, Tensor, Tensor]:
@@ -393,7 +394,17 @@ def _multiclass_stat_scores_update(
         fn = ((target_oh != preds_oh) & (target_oh == 1)).sum(sum_dim)
         fp = ((target_oh != preds_oh) & (target_oh == 0)).sum(sum_dim)
         tn = ((target_oh == preds_oh) & (target_oh == 0)).sum(sum_dim)
-        return tp, fp, tn, fn
+    elif average == "micro":
+        preds = preds.flatten()
+        target = target.flatten()
+        if ignore_index is not None:
+            idx = target != ignore_index
+            preds = preds[idx]
+            target = target[idx]
+        tp = (preds == target).sum()
+        fp = (preds != target).sum()
+        fn = (preds != target).sum()
+        tn = num_classes * preds.numel() - (fp + fn + tp)
     else:
         preds = preds.flatten()
         target = target.flatten()
@@ -408,7 +419,7 @@ def _multiclass_stat_scores_update(
         fp = confmat.sum(0) - tp
         fn = confmat.sum(1) - tp
         tn = confmat.sum() - (fp + fn + tp)
-        return tp, fp, tn, fn
+    return tp, fp, tn, fn
 
 
 def _multiclass_stat_scores_compute(
@@ -426,8 +437,8 @@ def _multiclass_stat_scores_compute(
     res = torch.stack([tp, fp, tn, fn, tp + fn], dim=-1)
     sum_dim = 0 if multidim_average == "global" else 1
     if average == "micro":
-        return res.sum(sum_dim)
-    elif average == "macro":
+        return res.sum(sum_dim) if res.ndim > 1 else res
+    if average == "macro":
         return res.float().mean(sum_dim)
     elif average == "weighted":
         weight = tp + fn
@@ -549,7 +560,9 @@ def multiclass_stat_scores(
         _multiclass_stat_scores_arg_validation(num_classes, top_k, average, multidim_average, ignore_index)
         _multiclass_stat_scores_tensor_validation(preds, target, num_classes, multidim_average, ignore_index)
     preds, target = _multiclass_stat_scores_format(preds, target, top_k)
-    tp, fp, tn, fn = _multiclass_stat_scores_update(preds, target, num_classes, top_k, multidim_average, ignore_index)
+    tp, fp, tn, fn = _multiclass_stat_scores_update(
+        preds, target, num_classes, top_k, average, multidim_average, ignore_index
+    )
     return _multiclass_stat_scores_compute(tp, fp, tn, fn, average, multidim_average)