Repository: Lightning-AI/torchmetrics

Issue:
AUROC thresholds are persistent
## 🐛 Bug?

When using metrics from the AUROC family (https://torchmetrics.readthedocs.io/en/stable/classification/auroc.html) and setting the argument `thresholds` to a value different than `None`, a _persistent_ buffer is created under the name `thresholds` storing the precomputed tensor of thresholds in memory. As far as I can see the persistence of this buffer is unnecessary, and it is problematic in various scenarios outlined below.

### To Reproduce

Consider the simplified PyTorch Lightning workflow below in which an AUROC metric is tracked during _testing_ only, since the metric is too expensive to track during training. As shown one must use the keyword argument `strict = False` in order to test a fitted model.

```python
import pytorch_lightning as pl
from torchmetrics.classification import BinaryAUROC


class LightningModel(pl.LightningModule):
    def __init__(self):
        super().__init__()

    def setup(self, stage):
        if stage == "test":
            self.auroc_test = BinaryAUROC(thresholds=42)


model_train = LightningModel()
model_train.setup("train")
state_dict = model_train.state_dict()

model_test_1 = LightningModel()
model_test_1.setup("test")
model_test_1.load_state_dict(state_dict.copy(), strict=False)  # runs fine

model_test_2 = LightningModel()
model_test_2.setup("test")
model_test_2.load_state_dict(state_dict.copy(), strict=True)  # crashes
```

### Why This is a Bug/Problem

1. Adding the `strict=False` keyword argument seems hacky.
2. Alternatively one can initialize test metrics during training, but this seems wasteful and overly complicated.
3. When using PyTorch Lightning one can add the `strict=False` keyword argument to the `load_from_checkpoint` method, but this is not an option when using default checkpoint saving combined with the `ckpt_path` argument in `Trainer` (as well as in `LightningCLI`).
4. From a purely conceptual perspective it makes sense that `thresholds` should be non-persistent since
a. they remain constant after initialization,
b. initialization is cheap.
See the link in the Additional Context section below for further explanation of this.

Item 3 can be resolved by requesting the PyTorch Lightning team to add a `strict` option to checkpoint loading in their `Trainer`, but to me that solution seems incomplete since the root of the issue seems to be that AUROC thresholds should be non-persistent, as outlined in item 4.

</details>

### Solution

Simply make AUROC thresholds non-persistent by replacing the line
```python
self.register_buffer("thresholds", thresholds)
```
by the line
```python
self.register_buffer("thresholds", thresholds, persistent=False)
```
in lines 148, 314, and 493 in the file `torchmetrics/classification/precision_recall_curve.py`.
__I have verified that this breaks no tests.__

### Additional Context
See this resolved GitHub issue regarding the motivation behind non-persistent buffers: https://github.com/pytorch/pytorch/issues/18056.

### Additional Remarks
There are a few other places in `torchmetrics` where buffers are used. It may be worthwhile to consider whether these should be persistent or not, and document and/or test this.


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

Resolving patch:
diff --git a/src/torchmetrics/classification/precision_recall_curve.py b/src/torchmetrics/classification/precision_recall_curve.py
--- a/src/torchmetrics/classification/precision_recall_curve.py
+++ b/src/torchmetrics/classification/precision_recall_curve.py
@@ -145,7 +145,7 @@ def __init__(
             self.add_state("preds", default=[], dist_reduce_fx="cat")
             self.add_state("target", default=[], dist_reduce_fx="cat")
         else:
-            self.register_buffer("thresholds", thresholds)
+            self.register_buffer("thresholds", thresholds, persistent=False)
             self.add_state(
                 "confmat", default=torch.zeros(len(thresholds), 2, 2, dtype=torch.long), dist_reduce_fx="sum"
             )
@@ -311,7 +311,7 @@ def __init__(
             self.add_state("preds", default=[], dist_reduce_fx="cat")
             self.add_state("target", default=[], dist_reduce_fx="cat")
         else:
-            self.register_buffer("thresholds", thresholds)
+            self.register_buffer("thresholds", thresholds, persistent=False)
             self.add_state(
                 "confmat",
                 default=torch.zeros(len(thresholds), num_classes, 2, 2, dtype=torch.long),
@@ -490,7 +490,7 @@ def __init__(
             self.add_state("preds", default=[], dist_reduce_fx="cat")
             self.add_state("target", default=[], dist_reduce_fx="cat")
         else:
-            self.register_buffer("thresholds", thresholds)
+            self.register_buffer("thresholds", thresholds, persistent=False)
             self.add_state(
                 "confmat",
                 default=torch.zeros(len(thresholds), num_labels, 2, 2, dtype=torch.long),
diff --git a/src/torchmetrics/functional/image/lpips.py b/src/torchmetrics/functional/image/lpips.py
--- a/src/torchmetrics/functional/image/lpips.py
+++ b/src/torchmetrics/functional/image/lpips.py
@@ -195,8 +195,8 @@ class ScalingLayer(nn.Module):
 
     def __init__(self) -> None:
         super().__init__()
-        self.register_buffer("shift", torch.Tensor([-0.030, -0.088, -0.188])[None, :, None, None])
-        self.register_buffer("scale", torch.Tensor([0.458, 0.448, 0.450])[None, :, None, None])
+        self.register_buffer("shift", torch.Tensor([-0.030, -0.088, -0.188])[None, :, None, None], persistent=False)
+        self.register_buffer("scale", torch.Tensor([0.458, 0.448, 0.450])[None, :, None, None], persistent=False)
 
     def forward(self, inp: Tensor) -> Tensor:
         """Process input."""
diff --git a/src/torchmetrics/metric.py b/src/torchmetrics/metric.py
--- a/src/torchmetrics/metric.py
+++ b/src/torchmetrics/metric.py
@@ -1024,12 +1024,12 @@ def __init__(
         self.op = operator
 
         if isinstance(metric_a, Tensor):
-            self.register_buffer("metric_a", metric_a)
+            self.register_buffer("metric_a", metric_a, persistent=False)
         else:
             self.metric_a = metric_a
 
         if isinstance(metric_b, Tensor):
-            self.register_buffer("metric_b", metric_b)
+            self.register_buffer("metric_b", metric_b, persistent=False)
         else:
             self.metric_b = metric_b
 


Test patch:
diff --git a/tests/unittests/classification/test_precision_recall_curve.py b/tests/unittests/classification/test_precision_recall_curve.py
--- a/tests/unittests/classification/test_precision_recall_curve.py
+++ b/tests/unittests/classification/test_precision_recall_curve.py
@@ -416,3 +416,18 @@ def test_valid_input_thresholds(metric, thresholds):
     with pytest.warns(None) as record:
         metric(thresholds=thresholds)
     assert len(record) == 0
+
+
+@pytest.mark.parametrize(
+    "metric",
+    [
+        BinaryPrecisionRecallCurve,
+        partial(MulticlassPrecisionRecallCurve, num_classes=NUM_CLASSES),
+        partial(MultilabelPrecisionRecallCurve, num_labels=NUM_CLASSES),
+    ],
+)
+@pytest.mark.parametrize("thresholds", [None, 100, [0.3, 0.5, 0.7, 0.9], torch.linspace(0, 1, 10)])
+def test_empty_state_dict(metric, thresholds):
+    """Test that metric have an empty state dict."""
+    m = metric(thresholds=thresholds)
+    assert m.state_dict() == {}, "Metric state dict should be empty."