Repository: Lightning-AI/torchmetrics

Issue:
Accuracy calculation blowing up memory use
## 🐛 Bug

<!-- A clear and concise description of what the bug is. -->

### To Reproduce

Steps to reproduce the behavior...


<!-- If you have a code sample, error messages, stack traces, please provide it here as well -->

<details>
  <summary>Code sample</summary>

```python
import torch
from torchmetrics import Accuracy

print(torch.cuda.memory_allocated()/1024**2)
# 0.0

device = "cuda"


x = torch.randn(8, 512, 30522, device=device)
preds = x.argmax(dim = -1)
y = torch.randint(0, 30522, (8, 512), device=device)
print(torch.cuda.memory_allocated()/1024**2)
# 476.96875

for i in range(100):
  gen_accuracy = Accuracy(task="multiclass", num_classes=30522, ignore_index=-100, average = 'micro', multidim_average = 'samplewise').to(device)
  acc1 = torch.sum((preds == y).float(), dim = -1) #doesn't blow up memory
  acc2 = gen_accuracy(preds, y) #does blow up memory
  print(torch.cuda.memory_allocated()/1024**2)
```


<!-- Ideally attach a minimal code sample to reproduce the decried issue.
Minimal means having the shortest code but still preserving the bug. -->

</details>

### Expected behavior

<!-- A clear and concise description of what you expected to happen. -->

Memory footprint stays constant and does not increase

### Environment

- TorchMetrics version (and how you installed TM, e.g. `conda`, `pip`, build from source):
- Python & PyTorch Version (e.g., 1.0):
- Any other relevant information such as OS (e.g., Linux):

Running this on google colab with a `!pip install torchmetrics` right before should be sufficient
### Additional context

<!-- Add any other context about the problem here. -->



Discussion:
Hi! thanks for your contribution!, great first issue!
additionally, this bug only seems to happen when multi-dim averaging is set to `samplewise`
@surya-narayanan thanks for reporting this issue. I try to figure out what is going on with the metric or if it is a general problem.
For now I can give you a stupid work around: simply add
```python
del gen_accuracy.tp, gen_accuracy.fp, gen_accuracy.tn, gen_accuracy.fn
```
to the end of your loop. That will fix the issue. 
This also indicates that it is some weird reference to metric states that are not being cleared as it should, and the memory is therefore not released at the end of the loop.

While tracking down what is going on here, I am also going to add that the comparison is not fair because for the code provided you are only calculating the value for the current batch and not utilizing the accumulation of the metric class at all. A better comparison would be with the functional version e.g.:
```python
for i in range(20):
    acc1 = torch.sum((preds == y).float(), dim = -1) #doesn't blow up memory
    acc2 = accuracy(preds, y, task="multiclass", num_classes=30522, ignore_index=-100, average = 'micro', multidim_average = 'samplewise') #doesn't blow up memory
    print(torch.cuda.memory_allocated()/1024**2)
```



Wow, thanks Nicki. I was wondeirng, what did you mean by "_not utilizing the accumulation of the metric class_"? Did you mean I wasn't doing something like 

```
import torch
from torchmetrics import Accuracy

print(torch.cuda.memory_allocated()/1024**2)
# 0.0

device = "cuda"


x = torch.randn(8, 512, 30522, device=device)
preds = x.argmax(dim = -1)
y = torch.randint(0, 30522, (8, 512), device=device)
print(torch.cuda.memory_allocated()/1024**2)
# 476.96875

myacc = 0

for i in range(100):
  gen_accuracy = Accuracy(task="multiclass", num_classes=30522, ignore_index=-100, average = 'micro', multidim_average = 'samplewise').to(device)
  acc1 = torch.sum((preds == y).float(), dim = -1) 
  acc2 = gen_accuracy(preds, y) 
  myacc += acc2 # is this what you meant by accumulation of the metric class?
  print(torch.cuda.memory_allocated()/1024**2)
```

Also, will assume that this is unresolved, but let me know if otherwise :) 

Resolving patch:
diff --git a/src/torchmetrics/metric.py b/src/torchmetrics/metric.py
--- a/src/torchmetrics/metric.py
+++ b/src/torchmetrics/metric.py
@@ -73,7 +73,8 @@ class Metric(Module, ABC):
                 custom implementation that calls ``torch.distributed.all_gather`` internally.
             - distributed_available_fn: function that checks if the distributed backend is available.
                 Defaults to a check of ``torch.distributed.is_available()`` and ``torch.distributed.is_initialized()``.
-            - sync_on_compute: If metric state should synchronize when ``compute`` is called. Default is ``True``-
+            - sync_on_compute: If metric state should synchronize when ``compute`` is called. Default is ``True``
+            - compute_with_cache: If results from ``compute`` should be cached. Default is ``False``
     """
 
     __jit_ignored_attributes__ = ["device"]
@@ -131,6 +132,11 @@ def __init__(
             raise ValueError(
                 f"Expected keyword argument `sync_on_compute` to be a `bool` but got {self.sync_on_compute}"
             )
+        self.compute_with_cache = kwargs.pop("compute_with_cache", True)
+        if not isinstance(self.compute_with_cache, bool):
+            raise ValueError(
+                f"Expected keyword argument `compute_with_cache` to be a `bool` but got {self.compute_with_cache}"
+            )
 
         if kwargs:
             kwargs_ = [f"`{a}`" for a in sorted(kwargs)]
@@ -559,10 +565,12 @@ def wrapped_func(*args: Any, **kwargs: Any) -> Any:
                 should_sync=self._to_sync,
                 should_unsync=self._should_unsync,
             ):
-                value = compute(*args, **kwargs)
-                self._computed = _squeeze_if_scalar(value)
+                value = _squeeze_if_scalar(compute(*args, **kwargs))
+
+            if self.compute_with_cache:
+                self._computed = value
 
-            return self._computed
+            return value
 
         return wrapped_func
 


Test patch:
diff --git a/tests/unittests/bases/test_metric.py b/tests/unittests/bases/test_metric.py
--- a/tests/unittests/bases/test_metric.py
+++ b/tests/unittests/bases/test_metric.py
@@ -45,6 +45,12 @@ def test_error_on_wrong_input():
     with pytest.raises(ValueError, match="Expected keyword argument `compute_on_cpu` to be an `bool` bu.*"):
         DummyMetric(compute_on_cpu=None)
 
+    with pytest.raises(ValueError, match="Expected keyword argument `sync_on_compute` to be a `bool` but.*"):
+        DummyMetric(sync_on_compute=None)
+
+    with pytest.raises(ValueError, match="Expected keyword argument `compute_with_cache` to be a `bool` but got.*"):
+        DummyMetric(compute_with_cache=None)
+
     with pytest.raises(ValueError, match="Unexpected keyword arguments: `foo`"):
         DummyMetric(foo=True)
 
@@ -151,7 +157,8 @@ def update(self, x):
     assert a._computed is None
 
 
-def test_compute():
+@pytest.mark.parametrize("compute_with_cache", [True, False])
+def test_compute(compute_with_cache):
     """Test that `compute` method works as expected."""
 
     class A(DummyMetric):
@@ -161,17 +168,17 @@ def update(self, x):
         def compute(self):
             return self.x
 
-    a = A()
+    a = A(compute_with_cache=compute_with_cache)
     assert a.compute() == 0
     assert a.x == 0
     a.update(1)
     assert a._computed is None
     assert a.compute() == 1
-    assert a._computed == 1
+    assert a._computed == 1 if compute_with_cache else a._computed is None
     a.update(2)
     assert a._computed is None
     assert a.compute() == 3
-    assert a._computed == 3
+    assert a._computed == 3 if compute_with_cache else a._computed is None
 
     # called without update, should return cached value
     a._computed = 5
@@ -455,6 +462,29 @@ def get_memory_usage():
         assert base_memory_level >= memory, "memory increased above base level"
 
 
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="Test requires GPU.")
+def test_constant_memory_on_repeat_init():
+    """Test that when initializing a metric multiple times the memory does not increase.
+
+    This only works for metrics with `compute_with_cache=False` as otherwise the cache will keep a refence that python
+    gc will not be able to collect and clean.
+    """
+
+    def mem():
+        return torch.cuda.memory_allocated() / 1024**2
+
+    x = torch.randn(10000).cuda()
+
+    for i in range(100):
+        m = DummyListMetric(compute_with_cache=False).cuda()
+        m(x)
+        if i == 0:
+            after_one_iter = mem()
+
+        # allow for 5% flucturation due to measuring
+        assert after_one_iter * 1.05 >= mem(), "memory increased too much above base level"
+
+
 @pytest.mark.skipif(not torch.cuda.is_available(), reason="test requires gpu")
 def test_specific_error_on_wrong_device():
     """Test that a specific error is raised if we detect input and metric are on different devices."""