diff --git a/src/torchmetrics/functional/classification/precision_recall_curve.py b/src/torchmetrics/functional/classification/precision_recall_curve.py index f6b09289692..1bba01b8b2d 100644 --- a/src/torchmetrics/functional/classification/precision_recall_curve.py +++ b/src/torchmetrics/functional/classification/precision_recall_curve.py @@ -189,6 +189,14 @@ def _binary_precision_recall_curve_format( return preds, target, thresholds +# The vectorized update paths materialise an intermediate holding +# ``preds.numel() * len(thresholds)`` elements, so their peak memory grows with the +# number of thresholds. The input-size checks below cannot see that on their own, so a +# large ``thresholds`` silently defeats the constant-memory behaviour that argument is +# documented to provide. Bound the intermediate as well as the input. +_MAX_VECTORIZED_ELEMENTS = 1_000_000 + + def _binary_precision_recall_curve_update( preds: Tensor, target: Tensor, @@ -202,7 +210,7 @@ def _binary_precision_recall_curve_update( """ if thresholds is None: return preds, target - if preds.numel() <= 50_000: + if preds.numel() <= 50_000 and preds.numel() * len(thresholds) <= _MAX_VECTORIZED_ELEMENTS: update_fn = _binary_precision_recall_curve_update_vectorized else: update_fn = _binary_precision_recall_curve_update_loop @@ -479,7 +487,7 @@ def _multiclass_precision_recall_curve_update( return preds, target if average == "micro": return _binary_precision_recall_curve_update(preds, target, thresholds) - if preds.numel() * num_classes <= 1_000_000: + if preds.numel() * num_classes <= 1_000_000 and preds.numel() * len(thresholds) <= _MAX_VECTORIZED_ELEMENTS: update_fn = _multiclass_precision_recall_curve_update_vectorized else: update_fn = _multiclass_precision_recall_curve_update_loop diff --git a/tests/unittests/classification/test_precision_recall_curve.py b/tests/unittests/classification/test_precision_recall_curve.py index e65f3a7be4a..89ba35ed633 100644 --- a/tests/unittests/classification/test_precision_recall_curve.py +++ b/tests/unittests/classification/test_precision_recall_curve.py @@ -13,6 +13,7 @@ # limitations under the License. from functools import partial +from unittest import mock import numpy as np import pytest @@ -28,6 +29,12 @@ PrecisionRecallCurve, ) from torchmetrics.functional.classification.precision_recall_curve import ( + _MAX_VECTORIZED_ELEMENTS, + _binary_precision_recall_curve_update_loop, + _binary_precision_recall_curve_update_vectorized, + _multiclass_precision_recall_curve_update, + _multiclass_precision_recall_curve_update_loop, + _multiclass_precision_recall_curve_update_vectorized, binary_precision_recall_curve, multiclass_precision_recall_curve, multilabel_precision_recall_curve, @@ -489,3 +496,51 @@ def test_precision_nan_when_no_preds_meet_threshold(thresholds): assert torch.isnan(precision_bins[mask]).all(), f"Precision not NaN for thresholds {thres[mask]}" assert torch.all(recall_bins[mask] == 0.0), f"Recall not zero for thresholds {thres[mask]}" + + +@pytest.mark.parametrize("num_thresholds", [5, 200]) +def test_multiclass_update_memory_does_not_scale_with_thresholds(num_thresholds): + """The vectorized path allocates ``preds.numel() * len(thresholds)`` elements. + + Passing ``thresholds`` is documented to bound memory, so the dispatch must take the + number of thresholds into account rather than looking at the input size alone. + See https://github.com/Lightning-AI/torchmetrics/issues/3299. + + """ + preds = torch.randn(60000, 3) + target = torch.randint(0, 3, (60000,)) + thresholds = torch.linspace(0, 1, num_thresholds) + + intermediate = preds.numel() * num_thresholds + expected_vectorized = intermediate <= _MAX_VECTORIZED_ELEMENTS + + with mock.patch( + "torchmetrics.functional.classification.precision_recall_curve" + "._multiclass_precision_recall_curve_update_vectorized", + wraps=_multiclass_precision_recall_curve_update_vectorized, + ) as vectorized: + _multiclass_precision_recall_curve_update(preds, target, 3, thresholds, average=None) + + assert vectorized.called is expected_vectorized + + +def test_multiclass_update_paths_agree(): + """Falling back to the loop is only safe if it returns the same state.""" + preds = torch.randn(2000, 3) + target = torch.randint(0, 3, (2000,)) + for num_thresholds in (3, 20, 200): + thresholds = torch.linspace(0, 1, num_thresholds) + vectorized = _multiclass_precision_recall_curve_update_vectorized(preds, target, 3, thresholds) + looped = _multiclass_precision_recall_curve_update_loop(preds, target, 3, thresholds) + assert torch.equal(vectorized, looped) + + +def test_binary_update_paths_agree(): + """Same invariant for the binary dispatch.""" + preds = torch.randn(2000) + target = torch.randint(0, 2, (2000,)) + for num_thresholds in (3, 20, 200): + thresholds = torch.linspace(0, 1, num_thresholds) + vectorized = _binary_precision_recall_curve_update_vectorized(preds, target, thresholds) + looped = _binary_precision_recall_curve_update_loop(preds, target, thresholds) + assert torch.equal(vectorized, looped)