binary_precision_recall_curve with thresholds: memory issues

[jogepari]

🐛 Bug

Using binary_precision_recall_curve with thresholds set results in huge memory consumption (and fast decreasing time savings)

To Reproduce

Notebook demo using memory_profiler

Preview

import torchmetrics
torchmetrics.__version__

'0.11.4'

import torch
from sklearn.metrics import precision_recall_curve
from torchmetrics.functional.classification import binary_precision_recall_curve

%load_ext memory_profiler

n_samples = 20_000_000
proba = torch.rand(n_samples)
labels = torch.randint(2, (n_samples, ))

sklearn for comparison

%time %memit precision_recall_curve(labels.numpy(), proba.numpy())

peak memory: 1371.58 MiB, increment: 733.07 MiB
CPU times: total: 4.11 s
Wall time: 6.65 s

No thresholds

%time %memit binary_precision_recall_curve(proba, labels)

peak memory: 1313.06 MiB, increment: 672.86 MiB
CPU times: total: 6.95 s
Wall time: 6.74 s

Thresholds with steps: 0.1, 0.05, 0.025

ths_defined = torch.arange(0.1, 1.00000001, step=0.1)
%time %memit binary_precision_recall_curve(proba, labels, thresholds=ths_defined)

peak memory: 5166.79 MiB, increment: 4522.11 MiB
CPU times: total: 4.17 s
Wall time: 2.09 s

ths_defined = torch.arange(0.1, 1.00000001, step=0.05)
%time %memit binary_precision_recall_curve(proba, labels, thresholds=ths_defined)

peak memory: 9196.10 MiB, increment: 8551.08 MiB
CPU times: total: 7.94 s
Wall time: 2.94 s

ths_defined = torch.arange(0.1, 1.00000001, step=0.025)
%time %memit binary_precision_recall_curve(proba, labels, thresholds=ths_defined)

peak memory: 17557.59 MiB, increment: 16912.54 MiB
CPU times: total: 15.4 s
Wall time: 5.59 s

Expected behavior

Being both faster and more memory efficient than without thresholds set.

Environment

• TorchMetrics version (and how you installed TM, e.g. conda, pip, build from source): 0.11.4, conda-forge
• Python & PyTorch Version (e.g., 1.0): python 3.10.10, pytorch 2.0.0
• Any other relevant information such as OS (e.g., Linux): Windows 11

[github-actions]

Hi! thanks for your contribution!, great first issue!

[SkafteNicki]

Hi @jogepari,
Happy to report that this have already been fixed in PR #1493. Here is a snippet running your code on master branch:

As you can see memory is now constant over a certain point.

To get the changes now, please install from master:
pip install https://github.com/Lightning-AI/metrics/archive/master.zip
or wait for the next release.

[jogepari]

Thank you, @SkafteNicki ! May I ask, seeing the fix was merged into master on Feb 23rd, how come it wasn't included in two releases since then?

Also, on your screenshot, calculation time with step=0.025 is already exceeding one without thresholds, it this normal behaviour?

[SkafteNicki]

@jogepari I can see its due to it being in the changed section https://github.com/Lightning-AI/torchmetrics/blob/master/CHANGELOG.md#changed and not the bugfix fixed section. Only linked PRs in the fixed section is included in bugfix releases.
We can argue if this is a bugfix, but the author of the PR saw it more as in improvement/change because the algorithm is actually working, it is just consuming a lot of memory.

Also, on your screenshot, calculation time with step=0.025 is already exceeding one without thresholds, it this normal behaviour?

So the consequence of lowering the memory requirements is that we need to use an alternative algorithm that is slower when using a lot of thresholds. So yes, I would say the results are expected.