Repository: Lightning-AI/torchmetrics

Issue:
DDP map evaluation with iou_type segm
When I use a single gpu, there is no error. Everything goes well.
When applying DDP, the following error occurs.

```
File "/data/kyanchen/miniconda3/envs/torch1mmcv2/lib/python3.10/site-packages/torchmetrics/detection/mean_ap.py", line 899, in compute
    self.lightning_module.cpu()
  File "/data/kyanchen/miniconda3/envs/torch1mmcv2/lib/python3.10/site-packages/lightning/fabric/utilities/device_dtype_mixin.py", line 78, in cpu
    return super().cpu()
  File "/data/kyanchen/miniconda3/envs/torch1mmcv2/lib/python3.10/site-packages/mmengine/model/base_model/base_model.py", line 246, in cpu
    return super().cpu()
  File "/data/kyanchen/miniconda3/envs/torch1mmcv2/lib/python3.10/site-packages/torch/nn/modules/module.py", line 798, in cpu
    precisions, recalls = self._calculate(classes)
  File "/data/kyanchen/miniconda3/envs/torch1mmcv2/lib/python3.10/site-packages/torchmetrics/detection/mean_ap.py", line 745, in _calculate
    return self._apply(lambda t: t.cpu())
  File "/data/kyanchen/miniconda3/envs/torch1mmcv2/lib/python3.10/site-packages/torch/nn/modules/module.py", line 641, in _apply
    ious = {
  File "/data/kyanchen/miniconda3/envs/torch1mmcv2/lib/python3.10/site-packages/torchmetrics/detection/mean_ap.py", line 746, in <dictcomp>
    module._apply(fn)
  File "/data/kyanchen/miniconda3/envs/torch1mmcv2/lib/python3.10/site-packages/torch/nn/modules/module.py", line 641, in _apply
    (idx, class_id): self._compute_iou(idx, class_id, max_detections)
  File "/data/kyanchen/miniconda3/envs/torch1mmcv2/lib/python3.10/site-packages/torchmetrics/detection/mean_ap.py", line 484, in _compute_iou
    module._apply(fn)
  File "/data/kyanchen/miniconda3/envs/torch1mmcv2/lib/python3.10/site-packages/torchmetrics/metric.py", line 651, in _apply
    gt = [gt[i] for i in gt_label_mask]
  File "/data/kyanchen/miniconda3/envs/torch1mmcv2/lib/python3.10/site-packages/torchmetrics/detection/mean_ap.py", line 484, in <listcomp>
    setattr(this, key, [fn(cur_v) for cur_v in current_val])
  File "/data/kyanchen/miniconda3/envs/torch1mmcv2/lib/python3.10/site-packages/torchmetrics/metric.py", line 651, in <listcomp>
    gt = [gt[i] for i in gt_label_mask]
IndexError: tuple index out of range


During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/mnt/search01/usr/chenkeyan/codes/lightning_framework/tools/train.py", line 57, in <module>
    setattr(this, key, [fn(cur_v) for cur_v in current_val])
  File "/data/kyanchen/miniconda3/envs/torch1mmcv2/lib/python3.10/site-packages/torch/nn/modules/module.py", line 798, in <lambda>
    return self._apply(lambda t: t.cpu())
AttributeError: 'tuple' object has no attribute 'cpu'
    main()
  File "/mnt/search01/usr/chenkeyan/codes/lightning_framework/tools/train.py", line 53, in main
    runner.run(args.status, ckpt_path=args.ckpt_path)
  File "/mnt/search01/usr/chenkeyan/codes/lightning_framework/tools/../mmpl/engine/runner/pl_runner.py", line 275, in run
    return trainer_func(model=self.model, datamodule=self.datamodule, *args, **kwargs)
  File "/data/kyanchen/miniconda3/envs/torch1mmcv2/lib/python3.10/site-packages/lightning/pytorch/trainer/trainer.py", line 520, in fit
    call._call_and_handle_interrupt(
  File "/data/kyanchen/miniconda3/envs/torch1mmcv2/lib/python3.10/site-packages/lightning/pytorch/trainer/call.py", line 68, in _call_and_handle_interrupt
    trainer._teardown()
  File "/data/kyanchen/miniconda3/envs/torch1mmcv2/lib/python3.10/site-packages/lightning/pytorch/trainer/trainer.py", line 958, in _teardown
    self.strategy.teardown()
  File "/data/kyanchen/miniconda3/envs/torch1mmcv2/lib/python3.10/site-packages/lightning/pytorch/strategies/ddp.py", line 430, in teardown
    super().teardown()
  File "/data/kyanchen/miniconda3/envs/torch1mmcv2/lib/python3.10/site-packages/lightning/pytorch/strategies/parallel.py", line 125, in teardown
    super().teardown()
  File "/data/kyanchen/miniconda3/envs/torch1mmcv2/lib/python3.10/site-packages/lightning/pytorch/strategies/strategy.py", line 475, in teardown
    self.lightning_module.cpu()
  File "/data/kyanchen/miniconda3/envs/torch1mmcv2/lib/python3.10/site-packages/lightning/fabric/utilities/device_dtype_mixin.py", line 78, in cpu
    return super().cpu()
  File "/data/kyanchen/miniconda3/envs/torch1mmcv2/lib/python3.10/site-packages/mmengine/model/base_model/base_model.py", line 246, in cpu
    return super().cpu()
  File "/data/kyanchen/miniconda3/envs/torch1mmcv2/lib/python3.10/site-packages/torch/nn/modules/module.py", line 798, in cpu
    return self._apply(lambda t: t.cpu())
  File "/data/kyanchen/miniconda3/envs/torch1mmcv2/lib/python3.10/site-packages/torch/nn/modules/module.py", line 641, in _apply
    module._apply(fn)
  File "/data/kyanchen/miniconda3/envs/torch1mmcv2/lib/python3.10/site-packages/torch/nn/modules/module.py", line 641, in _apply
    module._apply(fn)
  File "/data/kyanchen/miniconda3/envs/torch1mmcv2/lib/python3.10/site-packages/torchmetrics/metric.py", line 651, in _apply
    setattr(this, key, [fn(cur_v) for cur_v in current_val])
  File "/data/kyanchen/miniconda3/envs/torch1mmcv2/lib/python3.10/site-packages/torchmetrics/metric.py", line 651, in <listcomp>
    setattr(this, key, [fn(cur_v) for cur_v in current_val])
  File "/data/kyanchen/miniconda3/envs/torch1mmcv2/lib/python3.10/site-packages/torch/nn/modules/module.py", line 798, in <lambda>
    return self._apply(lambda t: t.cpu())
AttributeError: 'tuple' object has no attribute 'cpu'
```


Discussion:
Hi! thanks for your contribution!, great first issue!
Hi @KyanChen, thanks for reporting this issue.
Would it be possible for you to create a small reproduceable script that we can run to get the error?
I have found where the error is. Can the engineers of torchmetrics help the community solve the problem? That is the predictions and gts for segm IoUtype are tuple-based. Torchmetrics only allow list[tensor] based now, so it can not sync between different gpus when applying DDP mode.
> I have found where the error is. Can the engineers of torchmetrics help the community solve the problem? That is the predictions and gts for segm IoUtype are tuple-based. Torchmetrics only allow list[tensor] based now, so it can not sync between different gpus when applying DDP mode.

maybe you can try to change your predictions from tuple-based to list-based?
The operation is in the torchmetrics lib. It can not be changed outside. in "torchmetrics/detection/mean_ap.py"

    def _get_safe_item_values(self, item: Dict[str, Any]) -> Union[Tensor, Tuple]:
        if self.iou_type == "bbox":
            boxes = _fix_empty_tensors(item["boxes"])
            if boxes.numel() > 0:
                boxes = box_convert(boxes, in_fmt=self.box_format, out_fmt="xyxy")
            return boxes
        elif self.iou_type == "segm":
            masks = []

            for i in item["masks"].cpu().numpy():
                rle = mask_utils.encode(np.asfortranarray(i))
                masks.append((tuple(rle["size"]), rle["counts"]))

            return tuple(masks)
        else:
            raise Exception(f"IOU type {self.iou_type} is not supported")
I am pretty sure that I can come up with a fix, but I have a hard time reproducing the error.
Would it be possible to provide an example input that causes the error?
```
import lightning
import torch
import torch.nn as nn
from torchmetrics.detection import MeanAveragePrecision


class ToyNet(lightning.LightningModule):
    def __init__(self):
        super().__init__()
        self.linear = nn.Linear(1, 1)
        self.metric = MeanAveragePrecision(iou_type="segm")

    def training_step(self, batch, batch_idx):
        return None

    def validation_step(self, batch, batch_idx):
        preds = []
        targets = []
        for _ in range(2):
            result = dict()
            num_preds = torch.randint(0, 10, (1,)).item()
            # result['boxes'] = torch.rand((num_preds, 4), device=self.device)
            result['scores'] = torch.rand((num_preds,), device=self.device)
            result['labels'] = torch.randint(0, 10, (num_preds,),  device=self.device)
            result['masks'] = torch.randint(0, 2, (num_preds, 10, 10),  device=self.device).bool()
            preds.append(result)
            # parse gt
            gt = dict()
            num_gt = torch.randint(0, 10, (1,)).item()
            # gt['boxes'] = torch.rand((num_gt, 4), device=self.device)
            gt['labels'] = torch.randint(0, 10, (num_gt,),  device=self.device)
            gt['masks'] = torch.randint(0, 2, (num_gt, 10, 10),  device=self.device).bool()
            targets.append(gt)
        self.metric.update(preds, targets)

    def on_validation_epoch_end(self) -> None:
        self.log_dict(self.metric.compute(), prog_bar=True)

    def configure_optimizers(self):
        return torch.optim.Adam(self.parameters(), lr=0.02)


if __name__ == '__main__':
    model = ToyNet()
    trainer = lightning.Trainer(accelerator='gpu', devices=2, max_epochs=1)
    val_dataloader = torch.utils.data.DataLoader(torch.utils.data.TensorDataset(torch.rand((10, 1)), torch.rand((10, 1))))
    trainer.validate(model, val_dataloader)

```

Resolving patch:
diff --git a/src/torchmetrics/detection/mean_ap.py b/src/torchmetrics/detection/mean_ap.py
--- a/src/torchmetrics/detection/mean_ap.py
+++ b/src/torchmetrics/detection/mean_ap.py
@@ -12,10 +12,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import logging
-from typing import Any, Dict, List, Optional, Sequence, Tuple, Union
+from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple, Union
 
 import numpy as np
 import torch
+import torch.distributed as dist
 from torch import IntTensor, Tensor
 
 from torchmetrics.detection.helpers import _fix_empty_tensors, _input_validator
@@ -870,6 +871,48 @@ def compute(self) -> dict:
         metrics.classes = torch.tensor(classes, dtype=torch.int)
         return metrics
 
+    def _apply(self, fn: Callable) -> torch.nn.Module:
+        """Custom apply function.
+
+        Excludes the detections and groundtruths from the casting when the iou_type is set to `segm` as the state is
+        no longer a tensor but a tuple.
+        """
+        if self.iou_type == "segm":
+            this = super()._apply(fn, exclude_state=("detections", "groundtruths"))
+        else:
+            this = super()._apply(fn)
+        return this
+
+    def _sync_dist(self, dist_sync_fn: Optional[Callable] = None, process_group: Optional[Any] = None) -> None:
+        """Custom sync function.
+
+        For the iou_type `segm` the detections and groundtruths are no longer tensors but tuples. Therefore, we need
+        to gather the list of tuples and then convert it back to a list of tuples.
+
+        """
+        super()._sync_dist(dist_sync_fn=dist_sync_fn, process_group=process_group)
+
+        if self.iou_type == "segm":
+            self.detections = self._gather_tuple_list(self.detections, process_group)
+            self.groundtruths = self._gather_tuple_list(self.groundtruths, process_group)
+
+    @staticmethod
+    def _gather_tuple_list(list_to_gather: List[Tuple], process_group: Optional[Any] = None) -> List[Any]:
+        """Gather a list of tuples over multiple devices."""
+        world_size = dist.get_world_size(group=process_group)
+        list_gathered = [None] * world_size
+        dist.all_gather_object(list_gathered, list_to_gather, group=process_group)
+
+        for rank in range(1, world_size):
+            if list_gathered[rank] != list_gathered[0]:
+                raise ValueError(f"Rank {rank} and Rank 0 have different values for the list to gather.")
+        list_merged = []
+        for idx in range(len(list_gathered[0])):
+            for rank in range(world_size):
+                list_merged.append(list_gathered[rank][idx])
+
+        return list_merged
+
     def plot(
         self, val: Optional[Union[Dict[str, Tensor], Sequence[Dict[str, Tensor]]]] = None, ax: Optional[_AX_TYPE] = None
     ) -> _PLOT_OUT_TYPE:
diff --git a/src/torchmetrics/metric.py b/src/torchmetrics/metric.py
--- a/src/torchmetrics/metric.py
+++ b/src/torchmetrics/metric.py
@@ -703,11 +703,16 @@ def set_dtype(self, dst_type: Union[str, torch.dtype]) -> "Metric":
         out._dtype_convert = False
         return out
 
-    def _apply(self, fn: Callable) -> Module:
+    def _apply(self, fn: Callable, exclude_state: Sequence[str] = "") -> Module:
         """Overwrite _apply function such that we can also move metric states to the correct device.
 
         This method is called by the base ``nn.Module`` class whenever `.to`, `.cuda`, `.float`, `.half` etc. methods
         are called. Dtype conversion is garded and will only happen through the special `set_dtype` method.
+
+        Args:
+            fn: the function to apply
+            exclude_state: list of state variables to exclude from applying the function, that then needs to be handled
+                by the metric class itself.
         """
         this = super()._apply(fn)
         fs = str(fn)
@@ -717,6 +722,9 @@ def _apply(self, fn: Callable) -> Module:
 
         # Also apply fn to metric states and defaults
         for key, value in this._defaults.items():
+            if key in exclude_state:
+                continue
+
             if isinstance(value, Tensor):
                 this._defaults[key] = fn(value)
             elif isinstance(value, Sequence):


Test patch:
diff --git a/tests/unittests/detection/test_map.py b/tests/unittests/detection/test_map.py
--- a/tests/unittests/detection/test_map.py
+++ b/tests/unittests/detection/test_map.py
@@ -50,6 +50,18 @@ def _create_inputs_masks() -> Input:
                     "labels": IntTensor([3, 2]),
                 },  # 73
             ],
+            [
+                {
+                    "masks": _mask_unsqueeze_bool(inputs_json["preds"][0]),
+                    "scores": Tensor([0.236]),
+                    "labels": IntTensor([4]),
+                },
+                {
+                    "masks": _masks_stack_bool([inputs_json["preds"][1], inputs_json["preds"][2]]),
+                    "scores": Tensor([0.318, 0.726]),
+                    "labels": IntTensor([3, 2]),
+                },  # 73
+            ],
         ],
         target=[
             [
@@ -59,6 +71,13 @@ def _create_inputs_masks() -> Input:
                     "labels": IntTensor([2, 2]),
                 },  # 73
             ],
+            [
+                {"masks": _mask_unsqueeze_bool(inputs_json["targets"][0]), "labels": IntTensor([4])},  # 42
+                {
+                    "masks": _masks_stack_bool([inputs_json["targets"][1], inputs_json["targets"][2]]),
+                    "labels": IntTensor([2, 2]),
+                },  # 73
+            ],
         ],
     )
 
@@ -357,7 +376,7 @@ def test_map_bbox(self, compute_on_cpu, ddp):
             metric_args={"class_metrics": True, "compute_on_cpu": compute_on_cpu},
         )
 
-    @pytest.mark.parametrize("ddp", [False])
+    @pytest.mark.parametrize("ddp", [False, True])
     def test_map_segm(self, compute_on_cpu, ddp):
         """Test modular implementation for correctness."""
         _inputs_masks = _create_inputs_masks()
@@ -660,3 +679,40 @@ def test_error_on_wrong_input():
             [{"boxes": Tensor(), "scores": Tensor(), "labels": IntTensor()}],
             [{"boxes": Tensor(), "labels": []}],
         )
+
+
+def _generate_random_segm_input(device):
+    """Generate random inputs for mAP when iou_type=segm."""
+    preds = []
+    targets = []
+    for _ in range(2):
+        result = {}
+        num_preds = torch.randint(0, 10, (1,)).item()
+        result["scores"] = torch.rand((num_preds,), device=device)
+        result["labels"] = torch.randint(0, 10, (num_preds,), device=device)
+        result["masks"] = torch.randint(0, 2, (num_preds, 10, 10), device=device).bool()
+        preds.append(result)
+        gt = {}
+        num_gt = torch.randint(0, 10, (1,)).item()
+        gt["labels"] = torch.randint(0, 10, (num_gt,), device=device)
+        gt["masks"] = torch.randint(0, 2, (num_gt, 10, 10), device=device).bool()
+        targets.append(gt)
+    return preds, targets
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="test requires cuda")
+def test_device_changing():
+    """See issue: https://github.com/Lightning-AI/torchmetrics/issues/1743.
+
+    Checks that the custom apply function of the metric works as expected.
+    """
+    device = "cuda"
+    metric = MeanAveragePrecision(iou_type="segm").to(device)
+
+    for _ in range(2):
+        preds, targets = _generate_random_segm_input(device)
+        metric.update(preds, targets)
+
+    metric = metric.cpu()
+    val = metric.compute()
+    assert isinstance(val, dict)