ITADN

Multiclass accuracy requires num_classes**2 memory

#3343Openbarakugav 创建于 2026-03-21
bug / fixhelp wanted
B
barakugavcommented
## 🐛 Bug Not really a bug, but a limitation that can be solved easily. The current implementation for the `MulticlassAccuracy` uses quadratic memory relative to the number of classes. In my specific case i tried to use it for hundreds of thousands of classes, and got CUDA OOM. From looking at the code, looks like the full confusion matrix was created just for easier and cleaner code, but its easy to use only linear memory. Thanks in advance for the support ❤️ ### To Reproduce ```python import torch from torchmetrics.classification import MulticlassAccuracy metric = MulticlassAccuracy(average='macro', num_classes=1_000_000) metric.update(torch.tensor([1]), torch.tensor([1])) # <-- OOM here ``` ### Additional context Fix: ```diff # _multiclass_stat_scores_update - 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() - fp = confmat.sum(0) - tp - fn = confmat.sum(1) - tp - tn = confmat.sum() - (fp + fn + tp) + tp = _bincount(preds[target == preds], minlength=num_classes) + fp = _bincount(preds, minlength=num_classes) - tp + fn = _bincount(target, minlength=num_classes) - tp + tn = target.numel() - (tp + fp + fn) ``` https://github.com/Lightning-AI/torchmetrics/pull/3342
0 条评论