Repository: Lightning-AI/torchmetrics

Issue:
`MeanAveragePrecision` CPU overload when using multi-device ddp
## 🐛 Bug

<!-- A clear and concise description of what the bug is. -->

I was using `pytorchlightning.Trainer` in `ddp` mode with 8 GPUs on a HPC. There are 5k images in my validation set, and I find all the CPU cores are fully loaded (100%) for extremely long time (~2 hours) at `validation_epoch_end()` stage.

### To Reproduce

1. using `ddp` mode with `number of device > 2`
2. using large validation/test set.

#### Code sample

```python
class Model(pl.LightningModule):
    def __init__(self):
        self.det = FasterRCNN(...)
        self.map = MeanAveragePrecision(compute_on_step=False)
        ...

    ...

    def validation_step(self, batch, batch_idx):
        images, target = batch
        preds = self.det(images)
        self.map.update(preds=preds, target=target)

    def validation_epoch_end(self, outs):
        val_map_out = self.map.compute()  <-- CPU overload
        self.map.reset()
```

### Expected behavior

The same computation takes place in `N` separate sub-process simultaneously. If the computations are on GPU, they should finish in roughly the same time. But current implementation moves all tensors to CPU and CPU overload occurs while `N` increases.

- ddp with 1 GPU, number of fully loaded GPU cores 16/32, takes ~6 mins.
- ddp with 2 GPU, number of fully loaded GPU cores 32/32, takes ~6 mins.
- ddp with >2 GPU, number of fully loaded GPU cores 32/32, takes > 1 hours.

### Environment

- OS (e.g., Linux): Linux
- Python: 3.8
- PyTorch: 1.8.1 and 1.10.2
- pytorch-lightning: 1.5.10
- torchmetrics: 0.7.2

### Temp solution

My temp solution is re-enable the mAP computation on GPU, although the GPU version is slower than CPU version #677 

I commented the following lines and made other relevant variables on the correct device.

```python
    def compute(self) -> dict:
        # move everything to CPU, as we are faster here
        # self.detection_boxes = [box.cpu() for box in self.detection_boxes]
        # self.detection_labels = [label.cpu() for label in self.detection_labels]
        # self.detection_scores = [score.cpu() for score in self.detection_scores]
        # self.groundtruth_boxes = [box.cpu() for box in self.groundtruth_boxes]
        # self.groundtruth_labels = [label.cpu() for label in self.groundtruth_labels]
```

Should we make `compute_on_cpu` optional? and any suggestion according to TorchMetrics' API design?

I'm also keeping my eyes on the work of improving mAP performance on both CPU and GPU. #742 


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

Resolving patch:
diff --git a/torchmetrics/detection/mean_ap.py b/torchmetrics/detection/mean_ap.py
--- a/torchmetrics/detection/mean_ap.py
+++ b/torchmetrics/detection/mean_ap.py
@@ -482,10 +482,10 @@ def _evaluate_image(
         nb_iou_thrs = len(self.iou_thresholds)
         nb_gt = len(gt)
         nb_det = len(det)
-        gt_matches = torch.zeros((nb_iou_thrs, nb_gt), dtype=torch.bool)
-        det_matches = torch.zeros((nb_iou_thrs, nb_det), dtype=torch.bool)
+        gt_matches = torch.zeros((nb_iou_thrs, nb_gt), dtype=torch.bool, device=gt.device)
+        det_matches = torch.zeros((nb_iou_thrs, nb_det), dtype=torch.bool, device=gt.device)
         gt_ignore = ignore_area_sorted
-        det_ignore = torch.zeros((nb_iou_thrs, nb_det), dtype=torch.bool)
+        det_ignore = torch.zeros((nb_iou_thrs, nb_det), dtype=torch.bool, device=gt.device)
 
         if torch.numel(ious) > 0:
             for idx_iou, t in enumerate(self.iou_thresholds):
@@ -723,13 +723,13 @@ def __calculate_recall_precision_scores(
             recall[idx, idx_cls, idx_bbox_area, idx_max_det_thrs] = rc[-1] if nd else 0
 
             # Remove zigzags for AUC
-            diff_zero = torch.zeros((1,))
-            diff = torch.ones((1,))
+            diff_zero = torch.zeros((1,), device=pr.device)
+            diff = torch.ones((1,), device=pr.device)
             while not torch.all(diff == 0):
                 diff = torch.clamp(torch.cat((pr[1:] - pr[:-1], diff_zero), 0), min=0)
                 pr += diff
 
-            inds = torch.searchsorted(rc, rec_thresholds, right=False)
+            inds = torch.searchsorted(rc, rec_thresholds.to(rc.device), right=False)
             num_inds = inds.argmax() if inds.max() >= nd else nb_rec_thrs
             inds = inds[:num_inds]
             prec[:num_inds] = pr[inds]
@@ -766,14 +766,6 @@ def compute(self) -> dict:
             - map_per_class: ``torch.Tensor`` (-1 if class metrics are disabled)
             - mar_100_per_class: ``torch.Tensor`` (-1 if class metrics are disabled)
         """
-
-        # move everything to CPU, as we are faster here
-        self.detection_boxes = [box.cpu() for box in self.detection_boxes]
-        self.detection_labels = [label.cpu() for label in self.detection_labels]
-        self.detection_scores = [score.cpu() for score in self.detection_scores]
-        self.groundtruth_boxes = [box.cpu() for box in self.groundtruth_boxes]
-        self.groundtruth_labels = [label.cpu() for label in self.groundtruth_labels]
-
         classes = self._get_classes()
         precisions, recalls = self._calculate(classes)
         map_val, mar_val = self._summarize_results(precisions, recalls)
diff --git a/torchmetrics/metric.py b/torchmetrics/metric.py
--- a/torchmetrics/metric.py
+++ b/torchmetrics/metric.py
@@ -93,6 +93,8 @@ class Metric(Module, ABC):
 
         kwargs: additional keyword arguments, see :ref:`Metric kwargs` for more info.
 
+            - compute_on_cpu: If metric state should be stored on CPU during computations. Only works
+                for list states.
             - dist_sync_on_step: If metric state should synchronize on ``forward()``
             - process_group: The process group on which the synchronization is called
             - dist_sync_fn: function that performs the allgather option on the metric state
@@ -120,6 +122,11 @@ def __init__(
             warnings.warn(
                 "Argument `compute_on_step` is deprecated in v0.8 and will be removed in v0.9", DeprecationWarning
             )
+        self.compute_on_cpu = kwargs.pop("compute_on_cpu", False)
+        if not isinstance(self.compute_on_cpu, bool):
+            raise ValueError(
+                f"Expected keyword argument `compute_on_cpu` to be an `bool` but got {self.compute_on_cpu}"
+            )
 
         self.dist_sync_on_step = kwargs.pop("dist_sync_on_step", False)
         if not isinstance(self.dist_sync_on_step, bool):
@@ -243,6 +250,9 @@ def forward(self, *args: Any, **kwargs: Any) -> Any:
         self._to_sync = self.dist_sync_on_step  # type: ignore
         # skip restore cache operation from compute as cache is stored below.
         self._should_unsync = False
+        # skip computing on cpu for the batch
+        _temp_compute_on_cpu = self.compute_on_cpu
+        self.compute_on_cpu = False
 
         # save context before switch
         cache = {attr: getattr(self, attr) for attr in self._defaults}
@@ -262,6 +272,7 @@ def forward(self, *args: Any, **kwargs: Any) -> Any:
         self._to_sync = True
         self._computed = None
         self._enable_grad = False
+        self.compute_on_cpu = _temp_compute_on_cpu
 
         return self._forward_cache
 
@@ -294,14 +305,23 @@ def _sync_dist(self, dist_sync_fn: Callable = gather_all_tensors, process_group:
 
     def _wrap_update(self, update: Callable) -> Callable:
         @functools.wraps(update)
-        def wrapped_func(*args: Any, **kwargs: Any) -> Optional[Any]:
+        def wrapped_func(*args: Any, **kwargs: Any) -> None:
             self._computed = None
             self._update_called = True
             with torch.set_grad_enabled(self._enable_grad):
-                return update(*args, **kwargs)
+                update(*args, **kwargs)
+            if self.compute_on_cpu:
+                self._move_list_states_to_cpu()
 
         return wrapped_func
 
+    def _move_list_states_to_cpu(self) -> None:
+        """Move list states to cpu to save GPU memory."""
+        for key in self._defaults.keys():
+            current_val = getattr(self, key)
+            if isinstance(current_val, Sequence):
+                setattr(self, key, [cur_v.to("cpu") for cur_v in current_val])
+
     def sync(
         self,
         dist_sync_fn: Optional[Callable] = None,


Test patch:
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
@@ -37,6 +37,9 @@ def test_error_on_wrong_input():
     with pytest.raises(ValueError, match="Expected keyword argument `dist_sync_fn` to be an callable function.*"):
         DummyMetric(dist_sync_fn=[2, 3])
 
+    with pytest.raises(ValueError, match="Expected keyword argument `compute_on_cpu` to be an `bool` bu.*"):
+        DummyMetric(compute_on_cpu=None)
+
 
 def test_inherit():
     """Test that metric that inherits can be instanciated."""
diff --git a/tests/detection/test_map.py b/tests/detection/test_map.py
--- a/tests/detection/test_map.py
+++ b/tests/detection/test_map.py
@@ -164,6 +164,7 @@ def _compare_fn(preds, target) -> dict:
 
 
 @pytest.mark.skipif(_pytest_condition, reason="test requires that torchvision=>0.8.0 is installed")
+@pytest.mark.parametrize("compute_on_cpu", [True, False])
 class TestMAP(MetricTester):
     """Test the MAP metric for object detection predictions.
 
@@ -175,7 +176,7 @@ class TestMAP(MetricTester):
     atol = 1e-1
 
     @pytest.mark.parametrize("ddp", [False, True])
-    def test_map(self, ddp):
+    def test_map(self, compute_on_cpu, ddp):
         """Test modular implementation for correctness."""
         self.run_class_metric_test(
             ddp=ddp,
@@ -185,7 +186,7 @@ def test_map(self, ddp):
             sk_metric=_compare_fn,
             dist_sync_on_step=False,
             check_batch=False,
-            metric_args={"class_metrics": True},
+            metric_args={"class_metrics": True, "compute_on_cpu": compute_on_cpu},
         )
 
 
diff --git a/tests/image/test_inception.py b/tests/image/test_inception.py
--- a/tests/image/test_inception.py
+++ b/tests/image/test_inception.py
@@ -104,11 +104,12 @@ def __len__(self):
 
 @pytest.mark.skipif(not torch.cuda.is_available(), reason="test is too slow without gpu")
 @pytest.mark.skipif(not _TORCH_FIDELITY_AVAILABLE, reason="test requires torch-fidelity")
-def test_compare_is(tmpdir):
+@pytest.mark.parametrize("compute_on_cpu", [True, False])
+def test_compare_is(tmpdir, compute_on_cpu):
     """check that the hole pipeline give the same result as torch-fidelity."""
     from torch_fidelity import calculate_metrics
 
-    metric = InceptionScore(splits=1).cuda()
+    metric = InceptionScore(splits=1, compute_on_cpu=compute_on_cpu).cuda()
 
     # Generate some synthetic data
     img1 = torch.randint(0, 255, (100, 3, 299, 299), dtype=torch.uint8)
diff --git a/tests/regression/test_pearson.py b/tests/regression/test_pearson.py
--- a/tests/regression/test_pearson.py
+++ b/tests/regression/test_pearson.py
@@ -53,8 +53,9 @@ def _sk_pearsonr(preds, target):
 class TestPearsonCorrcoef(MetricTester):
     atol = 1e-2
 
+    @pytest.mark.parametrize("compute_on_cpu", [True, False])
     @pytest.mark.parametrize("ddp", [True, False])
-    def test_pearson_corrcoef(self, preds, target, ddp):
+    def test_pearson_corrcoef(self, preds, target, compute_on_cpu, ddp):
         self.run_class_metric_test(
             ddp=ddp,
             preds=preds,
@@ -62,6 +63,7 @@ def test_pearson_corrcoef(self, preds, target, ddp):
             metric_class=PearsonCorrCoef,
             sk_metric=_sk_pearsonr,
             dist_sync_on_step=False,
+            metric_args={"compute_on_cpu": compute_on_cpu},
         )
 
     def test_pearson_corrcoef_functional(self, preds, target):