Lightning-AI / Lightning-AI/torchmetrics

MulticlassPrecisionRecallCurve with thresholds=[integer] does not use constant memory on MPS

Open
#3,299 0 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

bug / fix help wanted
Dominant language
Python
Stars
2.5k
Forks
526
Avg merge
6d 11h
Merged PRs (30d)
5

Description

## 🐛 Bug

According to the [MulticlassPrecisionRecallCurve documentation](https://lightning.ai/docs/torchmetrics/stable/classification/precision_recall_curve.html#multiclassprecisionrecallcurve), setting the `thresholds` argument should cause MulticlassPrecisionRecallCurve to use a constant amount of memory.
When the metric is on MPS, this is not the case.

### To Reproduce

Sample code and steps to reproduce the behavior with expected result...

Code sample

```python
# Correct behavior
metric = torchmetrics.classification.MulticlassPrecisionRecallCurve(num_classes=3, ignore_index=-1, thresholds=200).to('cpu')
pred = torch.randn(60000, 3).to('cpu')
target = torch.randint(0, 3, (60000,)).to('cpu')
metric = metric(pred, target)
print(metric)

# Crashes with OOM error
metric = torchmetrics.classification.MulticlassPrecisionRecallCurve(num_classes=3, ignore_index=-1, thresholds=200).to('mps')
pred = torch.randn(60000, 3).to('mps')
target = torch.randint(0, 3, (60000,)).to('mps')
metric = metric(pred, target)
print(metric)
```

Environment

- TorchMetrics version (if build from source, add commit SHA): 1.8.1
- Python & PyTorch Version (e.g., 1.0): 3.11.7
- Any other relevant information such as OS (e.g., Linux): MacOS Tahoe 26.0.1 on Macbook Pro M4 Max

### Additional context

Here's the error I get when I run the code on MPS. I can see that it does not use constant memory, because it attempts to allocate different amounts of memory when I vary the input size.

Traceback

```
---------------------------------------------------------------------------
RuntimeError Traceback (most recent call last)
Cell In[15], [line 5](vscode-notebook-cell:?execution_count=15&line=5)
3 target = torch.randint(0, 3, (60000,)).to('mps')
4 print( pred.shape, target.shape)
----> [5](vscode-notebook-cell:?execution_count=15&line=5) metric = metric(pred, target)
6 metric

File ~/Programs/marcos/.venv/lib/python3.11/site-packages/torch/nn/modules/module.py:1773, in Module._wrapped_call_impl(self, *args, **kwargs)
1771 return self._compiled_call_impl(*args, **kwargs) # type: ignore[misc]
1772 else:
-> [1773](https://file+.vscode-resource.vscode-cdn.net/Users/pgq596/Programs/marcos/notebooks/~/Programs/marcos/.venv/lib/python3.11/site-packages/torch/nn/modules/module.py:1773) return self._call_impl(*args, **kwargs)

File ~/Programs/marcos/.venv/lib/python3.11/site-packages/torch/nn/modules/module.py:1784, in Module._call_impl(self, *args, **kwargs)
1779 # If we don't have any hooks, we want to skip the rest of the logic in
1780 # this function, and just call forward.
1781 if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks
1782 or _global_backward_pre_hooks or _global_backward_hooks
1783 or _global_forward_hooks or _global_forward_pre_hooks):
-> [1784](https://file+.vscode-resource.vscode-cdn.net/Users/pgq596/Programs/marcos/notebooks/~/Programs/marcos/.venv/lib/python3.11/site-packages/torch/nn/modules/module.py:1784) return forward_call(*args, **kwargs)
1786 result = None
1787 called_always_called_hooks = set()

File ~/Programs/marcos/.venv/lib/python3.11/site-packages/torchmetrics/metric.py:315, in Metric.forward(self, *args, **kwargs)
313 self._forward_cache = self._forward_full_state_update(*args, **kwargs)
314 else:
--> [315](https://file+.vscode-resource.vscode-cdn.net/Users/pgq596/Programs/marcos/notebooks/~/Programs/marcos/.venv/lib/python3.11/site-packages/torchmetrics/metric.py:315) self._forward_cache = self._forward_reduce_state_update(*args, **kwargs)
317 return self._forward_cache

File ~/Programs/marcos/.venv/lib/python3.11/site-packages/torchmetrics/metric.py:384, in Metric._forward_reduce_state_update(self, *args, **kwargs)
381 self._enable_grad = True # allow grads for batch computation
383 # calculate batch state and compute batch value
--> [384](https://file+.vscode-resource.vscode-cdn.net/Users/pgq596/Programs/marcos/notebooks/~/Programs/marcos/.venv/lib/python3.11/site-packages/torchmetrics/metric.py:384) self.update(*args, **kwargs)
385 batch_val = self.compute()
387 # reduce batch and global state

File ~/Programs/marcos/.venv/lib/python3.11/site-packages/torchmetrics/metric.py:559, in Metric._wrap_update..wrapped_func(*args, **kwargs)
551 if "Expected all tensors to be on" in str(err):
552 raise RuntimeError(
553 "Encountered different devices in metric calculation (see stacktrace for details)."
554 " This could be due to the metric class not being on the same device as input."
(...) 557 " device corresponds to the device of the input."
558 ) from err
--> [559](https://file+.vscode-resource.vscode-cdn.net/Users/pgq596/Programs/marcos/notebooks/~/Programs/marcos/.venv/lib/python3.11/site-packages/torchmetrics/metric.py:559) raise err
561 if self.compute_on_cpu:
562 self._move_list_states_to_cpu()

File ~/Programs/marcos/.venv/lib/python3.11/site-packages/torchmetrics/metric.py:549, in Metric._wrap_update..wrapped_func(*args, **kwargs)
547 with torch.set_grad_enabled(self._enable_grad):
548 try:
--> [549](https://file+.vscode-resource.vscode-cdn.net/Users/pgq596/Programs/marcos/notebooks/~/Programs/marcos/.venv/lib/python3.11/site-packages/torchmetrics/metric.py:549) update(*args, **kwargs)
550 except RuntimeError as err:
551 if "Expected all tensors to be on" in str(err):

File ~/Programs/marcos/.venv/lib/python3.11/site-packages/torchmetrics/classification/precision_recall_curve.py:379, in MulticlassPrecisionRecallCurve.update(self, preds, target)
375 _multiclass_precision_recall_curve_tensor_validation(preds, target, self.num_classes, self.ignore_index)
376 preds, target, _ = _multiclass_precision_recall_curve_format(
377 preds, target, self.num_classes, self.thresholds, self.ignore_index, self.average
378 )
--> [379](https://file+.vscode-resource.vscode-cdn.net/Users/pgq596/Programs/marcos/notebooks/~/Programs/marcos/.venv/lib/python3.11/site-packages/torchmetrics/classification/precision_recall_curve.py:379) state = _multiclass_precision_recall_curve_update(
380 preds, target, self.num_classes, self.thresholds, self.average
381 )
382 if isinstance(state, Tensor):
383 self.confmat += state

File ~/Programs/marcos/.venv/lib/python3.11/site-packages/torchmetrics/functional/classification/precision_recall_curve.py:486, in _multiclass_precision_recall_curve_update(preds, target, num_classes, thresholds, average)
484 else:
485 update_fn = _multiclass_precision_recall_curve_update_loop
--> [486](https://file+.vscode-resource.vscode-cdn.net/Users/pgq596/Programs/marcos/notebooks/~/Programs/marcos/.venv/lib/python3.11/site-packages/torchmetrics/functional/classification/precision_recall_curve.py:486) return update_fn(preds, target, num_classes, thresholds)

File ~/Programs/marcos/.venv/lib/python3.11/site-packages/torchmetrics/functional/classification/precision_recall_curve.py:507, in _multiclass_precision_recall_curve_update_vectorized(preds, target, num_classes, thresholds)
505 unique_mapping += 4 * torch.arange(num_classes, device=preds.device).unsqueeze(0).unsqueeze(-1)
506 unique_mapping += 4 * num_classes * torch.arange(len_t, device=preds.device)
--> [507](https://file+.vscode-resource.vscode-cdn.net/Users/pgq596/Programs/marcos/notebooks/~/Programs/marcos/.venv/lib/python3.11/site-packages/torchmetrics/functional/classification/precision_recall_curve.py:507) bins = _bincount(unique_mapping.flatten(), minlength=4 * num_classes * len_t)
508 return bins.reshape(len_t, num_classes, 2, 2)

File ~/Programs/marcos/.venv/lib/python3.11/site-packages/torchmetrics/utilities/data.py:203, in _bincount(x, minlength)
200 minlength = len(torch.unique(x))
202 if torch.are_deterministic_algorithms_enabled() or _XLA_AVAILABLE or x.is_mps:
--> [203](https://file+.vscode-resource.vscode-cdn.net/Users/pgq596/Programs/marcos/notebooks/~/Programs/marcos/.venv/lib/python3.11/site-packages/torchmetrics/utilities/data.py:203) mesh = torch.arange(minlength, device=x.device).repeat(len(x), 1)
204 return torch.eq(x.reshape(-1, 1), mesh).sum(dim=0)
206 return torch.bincount(x, minlength=minlength)

RuntimeError: Invalid buffer size: 643.73 GiB
```

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start with torchmetrics/classification/precision_recall_curve.py and the _multiclass_precision_recall_curve_update_vectorized and _bincount paths identified in the traceback. Compare the CPU and MPS behavior for integer thresholds using the supplied reproduction, then verify that MPS memory remains constant as the input size changes without the reported allocation failure.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
machine-learning
Issue type
Bug
Difficulty
3/5
Estimated time
1-2 days
Activity status
Stale
Clarity
Clearly specified
Newbie friendliness
55/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.