Repository: Lightning-AI/torchmetrics

Issue:
Inefficient memory consumption by MeanMetric
https://github.com/PyTorchLightning/metrics/blob/3af729508289d3babf0e166d9e8405cb2b0758a2/torchmetrics/aggregation.py#L333-L408

Old average meters eats nothing compare to 16+ gb RAB by MeanMetric
```py
class AverageMeter(object):
    """Computes and stores the average and current value"""

    def __init__(self):
        self.reset()

    def reset(self):
        self.val = .0
        self.avg = .0
        self.sum = .0
        self.count = .0

    def update(self, val, n=1):
        self.val = val
        self.sum += val * n
        self.count += n
        self.avg = self.sum / self.count
```


Discussion:
Hi @zuenko, could you please provide an example of such that we can debug?

Resolving patch:
diff --git a/torchmetrics/metric.py b/torchmetrics/metric.py
--- a/torchmetrics/metric.py
+++ b/torchmetrics/metric.py
@@ -144,6 +144,7 @@ def __init__(
         self._update_called = False
         self._to_sync = True
         self._should_unsync = True
+        self._enable_grad = False
 
         # initialize state
         self._defaults: Dict[str, Union[List, Tensor]] = {}
@@ -236,8 +237,8 @@ def forward(self, *args: Any, **kwargs: Any) -> Any:
                 "HINT: Did you forget to call ``unsync`` ?."
             )
 
-        with torch.no_grad():
-            self.update(*args, **kwargs)
+        # global accumulation
+        self.update(*args, **kwargs)
 
         self._to_sync = self.dist_sync_on_step  # type: ignore
         # skip restore cache operation from compute as cache is stored below.
@@ -247,6 +248,7 @@ def forward(self, *args: Any, **kwargs: Any) -> Any:
         cache = {attr: getattr(self, attr) for attr in self._defaults}
 
         # call reset, update, compute, on single batch
+        self._enable_grad = True  # allow grads for batch computation
         self.reset()
         self.update(*args, **kwargs)
         self._forward_cache = self.compute()
@@ -259,6 +261,7 @@ def forward(self, *args: Any, **kwargs: Any) -> Any:
         self._should_unsync = True
         self._to_sync = True
         self._computed = None
+        self._enable_grad = False
 
         return self._forward_cache
 
@@ -294,7 +297,8 @@ def _wrap_update(self, update: Callable) -> Callable:
         def wrapped_func(*args: Any, **kwargs: Any) -> Optional[Any]:
             self._computed = None
             self._update_called = True
-            return update(*args, **kwargs)
+            with torch.set_grad_enabled(self._enable_grad):
+                return update(*args, **kwargs)
 
         return wrapped_func
 


Test patch:
diff --git a/requirements/test.txt b/requirements/test.txt
--- a/requirements/test.txt
+++ b/requirements/test.txt
@@ -10,6 +10,7 @@ mypy>=0.790
 phmdoctest>=1.1.1
 pre-commit>=1.0
 
+psutil
 requests
 fire
 
diff --git a/tests/bases/test_metric.py b/tests/bases/test_metric.py
--- a/tests/bases/test_metric.py
+++ b/tests/bases/test_metric.py
@@ -11,11 +11,13 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+import os
 import pickle
 from collections import OrderedDict
 
 import cloudpickle
 import numpy as np
+import psutil
 import pytest
 import torch
 from torch import Tensor, nn, tensor
@@ -362,3 +364,44 @@ def device(self):
     assert module.device == module.metric.device
     if isinstance(module.metric.x, Tensor):
         assert module.device == module.metric.x.device
+
+
+@pytest.mark.parametrize("device", ["cpu", "cuda"])
+@pytest.mark.parametrize("requires_grad", [True, False])
+def test_constant_memory(device, requires_grad):
+    """Checks that when updating a metric the memory does not increase."""
+    if not torch.cuda.is_available() and device == "cuda":
+        pytest.skip("Test requires GPU support")
+
+    def get_memory_usage():
+        if device == "cpu":
+            pid = os.getpid()
+            py = psutil.Process(pid)
+            return py.memory_info()[0] / 2.0 ** 30
+        else:
+            return torch.cuda.memory_allocated()
+
+    x = torch.randn(10, requires_grad=requires_grad, device=device)
+
+    # try update method
+    metric = DummyMetricSum().to(device)
+
+    metric.update(x.sum())
+
+    # we allow for 5% flucturation due to measuring
+    base_memory_level = 1.05 * get_memory_usage()
+
+    for _ in range(10):
+        metric.update(x.sum())
+        memory = get_memory_usage()
+        assert base_memory_level >= memory, "memory increased above base level"
+
+    # try forward method
+    metric = DummyMetricSum().to(device)
+    metric(x.sum())
+    base_memory_level = get_memory_usage()
+
+    for _ in range(10):
+        metric.update(x.sum())
+        memory = get_memory_usage()
+        assert base_memory_level >= memory, "memory increased above base level"
diff --git a/tests/retrieval/helpers.py b/tests/retrieval/helpers.py
--- a/tests/retrieval/helpers.py
+++ b/tests/retrieval/helpers.py
@@ -485,7 +485,7 @@ def run_precision_test_gpu(
         metric_functional: Callable,
     ):
         if not torch.cuda.is_available():
-            pytest.skip()
+            pytest.skip("Test requires GPU")
 
         def metric_functional_ignore_indexes(preds, target, indexes):
             return metric_functional(preds, target)