Repository: Lightning-AI/torchmetrics

Issue:
Binned implementation of PrecisionRecallCurve is extremely slow for large batch sizes
## 🐛 Bug

Since the classification metrics refactor in torchmetrics 0.10, the update step of `PrecisionRecallCurve` and metrics based on it (such as `AveragePrecision`) is extremely slow and memory-consuming when using the binned implementation with many samples. This is, for example, the case in semantic segmentation scenarios, where every pixel is a sample.

The binned implementation is intended to be both faster and memory efficient than computing the exact precision-recall curve. This is currently not the case.

### To Reproduce

```python
import torch
from torchmetrics.classification import BinaryAveragePrecision

targets = torch.randint(0, 2, (16, 512, 512), device="cuda")
preds = torch.rand((16, 512, 512), device="cuda")

metric_exact = BinaryAveragePrecision().cuda()
metric_binned = BinaryAveragePrecision(thresholds=200).cuda()

with torch.inference_mode():
    start_exact_evt = torch.cuda.Event(enable_timing=True)
    update_exact_evt = torch.cuda.Event(enable_timing=True)
    compute_exact_evt = torch.cuda.Event(enable_timing=True)
    start_binned_evt = torch.cuda.Event(enable_timing=True)
    update_binned_evt = torch.cuda.Event(enable_timing=True)
    compute_binned_evt = torch.cuda.Event(enable_timing=True)
    
    start_exact_evt.record()
    metric_exact.update(preds, targets)
    update_exact_evt.record()
    ap_exact = metric_exact.compute()
    compute_exact_evt.record()
    
    start_binned_evt.record()
    metric_binned.update(preds, targets)
    update_binned_evt.record()
    ap_binned = metric_exact.compute()
    compute_binned_evt.record()

torch.cuda.synchronize()

print("AP (exact):", ap_exact)
print("AP (binned):", ap_binned)
print("Time (exact): {} ms + {} ms".format(
    start_exact_evt.elapsed_time(update_exact_evt),
    update_exact_evt.elapsed_time(compute_exact_evt)
))
print("Time (binned): {} ms + {} ms".format(
    start_binned_evt.elapsed_time(update_binned_evt),
    update_binned_evt.elapsed_time(compute_binned_evt)
))
```

Output (using an RTX 3090):
```
AP (exact): tensor(0.5008, device='cuda:0')
AP (binned): tensor(0.5008, device='cuda:0')
Time (exact): 1.8867520093917847 ms + 4.845024108886719 ms
Time (binned): 10816.78515625 ms + 0.002047999994829297 ms
```

Binned average precision was 1600x slower than exact average precision.
The memory consumption of exact AP on the GPU was 260 MB, while binned AP consumed 20,000 MB.

### Expected behavior

The binned implementation of `PrecisionRecallCurve`, i.e., with `thresholds != None`, should be faster and more memory-efficient.

### Environment

- TorchMetrics version (and how you installed TM, e.g. `conda`, `pip`, build from source): 0.11.1 (pip)
- Python & PyTorch Version (e.g., 1.0): Python 3.7.9, PyTorch 1.12.1
- Any other relevant information such as OS (e.g., Linux): Linux

### Additional context

The [current implementation of `_binary_precision_recall_curve_update`](https://github.com/Lightning-AI/metrics/blob/b95d482b81a78086764ee89cde0c211eb67dec4c/src/torchmetrics/functional/classification/precision_recall_curve.py#L185) compares all predictions with all thresholds at once, which is very memory consuming:

```python
len_t = len(thresholds)
preds_t = (preds.unsqueeze(-1) >= thresholds.unsqueeze(0)).long()  # num_samples x num_thresholds
unique_mapping = preds_t + 2 * target.unsqueeze(-1) + 4 * torch.arange(len_t, device=target.device)
bins = _bincount(unique_mapping.flatten(), minlength=4 * len_t)
return bins.reshape(len_t, 2, 2)
```

Computing `preds_t` and `unique_mapping` is quite fast, but both consume large amounts of memory (number of samples * number of thresholds * 8 byte). `_bincount`, on the other hand, does not consume much memory, but a lot of time.

Iterating over the thresholds one by one is significantly faster (0.6 sec instead of 10 sec) and only consumes 150 MB memory:

```python
confmat = thresholds.new_empty((len(thresholds), 2, 2), dtype=torch.int64)
for i in range(len(thresholds)):
    preds_t = (preds >= thresholds[i]).long()
    unique_mapping = preds_t + 2 * targets
    bins = torch.bincount(unique_mapping, minlength=4)
    confmat[i] = bins.reshape(2, 2)
```

Using the [implementation of `BinnedPrecisionRecallCurve` from torchmetrics 0.9](https://github.com/Lightning-AI/metrics/blob/ff61c482e5157b43e647565fa0020a4ead6e9d61/torchmetrics/classification/binned_precision_recall.py#L145) is even faster (0.1 sec) and uses only 60 MB:

```python
confmat = thresholds.new_empty((len(thresholds), 2, 2), dtype=torch.int64)
targets_t = targets == 1
for i in range(len(thresholds)):
    preds_t = preds >= thresholds[i]
    confmat[i, 1, 1] = (targets_t & preds_t).sum()
    confmat[i, 0, 1] = ((~targets_t) & preds_t).sum()
    confmat[i, 1, 0] = (targets_t & (~preds_t)).sum()
    confmat[i, 0, 0] = len(preds_t) - confmat[i, 0, 1] - confmat[i, 1, 0] - confmat[i, 1, 1]
```


Discussion:
I conducted some further experiments with respect to the number of samples, comparing the implementation of the BinaryAveragePrecision update step from torchmetrics 0.11 with that of torchmetrics 0.9:

|  #Samples  | 0.11 Speed | 0.9 Speed | 0.11 Memory | 0.9 Memory |
|-----------:|-----------:|----------:|------------:|-----------:|
|        100 |     0.2 ms |   23.8 ms |      < 1 MB |     < 1 MB |
|      1,000 |     1.1 ms |   24.4 ms |       20 MB |     < 1 MB |
|     25,000 |    24.0 ms |   25.8 ms |      200 MB |     < 1 MB |
|    100,000 |   115.3 ms |   25.5 ms |      722 MB |       2 MB |
|  1,000,000 | 1,012.1 ms |   26.7 ms |    7,156 MB |       2 MB |

For small numbers of samples, the new implementation is in fact faster, but never more memory-efficient.
However, it's time consumption scales linearly with the number of samples, while the old implementation is almost constant with respect to the number of samples. It will scale with the number of thresholds though (I used 200 here).  This makes it difficult to switch between implementations depending on the number of samples.

I personally would be in favor of switching back to the old implementation. It is not the fastest in all scenarios, but the most memory efficient and does not get terribly slow for large numbers of samples.

Resolving patch:
diff --git a/src/torchmetrics/functional/classification/precision_recall_curve.py b/src/torchmetrics/functional/classification/precision_recall_curve.py
--- a/src/torchmetrics/functional/classification/precision_recall_curve.py
+++ b/src/torchmetrics/functional/classification/precision_recall_curve.py
@@ -194,6 +194,23 @@ def _binary_precision_recall_curve_update(
     """
     if thresholds is None:
         return preds, target
+    if preds.numel() <= 50_000:
+        update_fn = _binary_precision_recall_curve_update_vectorized
+    else:
+        update_fn = _binary_precision_recall_curve_update_loop
+    return update_fn(preds, target, thresholds)
+
+
+def _binary_precision_recall_curve_update_vectorized(
+    preds: Tensor,
+    target: Tensor,
+    thresholds: Tensor,
+) -> Union[Tensor, Tuple[Tensor, Tensor]]:
+    """Returns the multi-threshold confusion matrix to calculate the pr-curve with.
+
+    This implementation is vectorized and faster than `_binary_precision_recall_curve_update_loop` for small
+    numbers of samples (up to 50k) but less memory- and time-efficient for more samples.
+    """
     len_t = len(thresholds)
     preds_t = (preds.unsqueeze(-1) >= thresholds.unsqueeze(0)).long()  # num_samples x num_thresholds
     unique_mapping = preds_t + 2 * target.unsqueeze(-1) + 4 * torch.arange(len_t, device=target.device)
@@ -201,6 +218,30 @@ def _binary_precision_recall_curve_update(
     return bins.reshape(len_t, 2, 2)
 
 
+def _binary_precision_recall_curve_update_loop(
+    preds: Tensor,
+    target: Tensor,
+    thresholds: Tensor,
+) -> Union[Tensor, Tuple[Tensor, Tensor]]:
+    """Returns the multi-threshold confusion matrix to calculate the pr-curve with.
+
+    This implementation loops over thresholds and is more memory-efficient than
+    `_binary_precision_recall_curve_update_vectorized`. However, it is slowwer for small
+    numbers of samples (up to 50k).
+    """
+    len_t = len(thresholds)
+    target = target == 1
+    confmat = thresholds.new_empty((len_t, 2, 2), dtype=torch.int64)
+    # Iterate one threshold at a time to conserve memory
+    for i in range(len_t):
+        preds_t = preds >= thresholds[i]
+        confmat[i, 1, 1] = (target & preds_t).sum()
+        confmat[i, 0, 1] = ((~target) & preds_t).sum()
+        confmat[i, 1, 0] = (target & (~preds_t)).sum()
+    confmat[:, 0, 0] = len(preds_t) - confmat[:, 0, 1] - confmat[:, 1, 0] - confmat[:, 1, 1]
+    return confmat
+
+
 def _binary_precision_recall_curve_compute(
     state: Union[Tensor, Tuple[Tensor, Tensor]],
     thresholds: Optional[Tensor],
@@ -409,8 +450,25 @@ def _multiclass_precision_recall_curve_update(
     """
     if thresholds is None:
         return preds, target
+    if preds.numel() * num_classes <= 1_000_000:
+        update_fn = _multiclass_precision_recall_curve_update_vectorized
+    else:
+        update_fn = _multiclass_precision_recall_curve_update_loop
+    return update_fn(preds, target, num_classes, thresholds)
+
+
+def _multiclass_precision_recall_curve_update_vectorized(
+    preds: Tensor,
+    target: Tensor,
+    num_classes: int,
+    thresholds: Tensor,
+) -> Union[Tensor, Tuple[Tensor, Tensor]]:
+    """Returns the multi-threshold confusion matrix to calculate the pr-curve with.
+
+    This implementation is vectorized and faster than `_binary_precision_recall_curve_update_loop` for small
+    numbers of samples but less memory- and time-efficient for more samples.
+    """
     len_t = len(thresholds)
-    # num_samples x num_classes x num_thresholds
     preds_t = (preds.unsqueeze(-1) >= thresholds.unsqueeze(0).unsqueeze(0)).long()
     target_t = torch.nn.functional.one_hot(target, num_classes=num_classes)
     unique_mapping = preds_t + 2 * target_t.unsqueeze(-1)
@@ -420,6 +478,31 @@ def _multiclass_precision_recall_curve_update(
     return bins.reshape(len_t, num_classes, 2, 2)
 
 
+def _multiclass_precision_recall_curve_update_loop(
+    preds: Tensor,
+    target: Tensor,
+    num_classes: int,
+    thresholds: Tensor,
+) -> Union[Tensor, Tuple[Tensor, Tensor]]:
+    """Returns the state to calculate the pr-curve with.
+
+    This implementation loops over thresholds and is more memory-efficient than
+    `_binary_precision_recall_curve_update_vectorized`. However, it is slowwer for small
+    numbers of samples.
+    """
+    len_t = len(thresholds)
+    target_t = torch.nn.functional.one_hot(target, num_classes=num_classes)
+    confmat = thresholds.new_empty((len_t, num_classes, 2, 2), dtype=torch.int64)
+    # Iterate one threshold at a time to conserve memory
+    for i in range(len_t):
+        preds_t = preds >= thresholds[i]
+        confmat[i, :, 1, 1] = (target_t & preds_t).sum(dim=0)
+        confmat[i, :, 0, 1] = ((~target_t) & preds_t).sum(dim=0)
+        confmat[i, :, 1, 0] = (target_t & (~preds_t)).sum(dim=0)
+    confmat[:, :, 0, 0] = len(preds_t) - confmat[:, :, 0, 1] - confmat[:, :, 1, 0] - confmat[:, :, 1, 1]
+    return confmat
+
+
 def _multiclass_precision_recall_curve_compute(
     state: Union[Tensor, Tuple[Tensor, Tensor]],
     num_classes: int,


Test patch: