COCO mAP metric

docs/source/metrics.rst

@@ -329,11 +329,13 @@ Complete list of metrics
     Frequency
     Loss
     MeanAbsoluteError
+    MeanAveragePrecision
     MeanPairwiseDistance
     MeanSquaredError
     metric.Metric
     metrics_lambda.MetricsLambda
     MultiLabelConfusionMatrix
+    ObjectDetectionMAP
     precision.Precision
     PSNR
     recall.Recall

ignite/distributed/utils.py

@@ -1,7 +1,8 @@
+import itertools
 import socket
 from contextlib import contextmanager
 from functools import wraps
-from typing import Any, Callable, List, Mapping, Optional, Tuple, Union
+from typing import Any, Callable, List, Mapping, Optional, Sequence, Tuple, Union
 
 import torch
 
@@ -350,29 +351,80 @@ def all_reduce(
     return _model.all_reduce(tensor, op, group=group)
 
 
+def _all_gather_tensors_with_shapes(
+    tensor: torch.Tensor, shapes: Sequence[Sequence[int]], group: Optional[Union[Any, List[int]]] = None
+) -> List[torch.Tensor]:
+    if _need_to_sync and isinstance(_model, _SerialModel):
+        sync(temporary=True)
+
+    if isinstance(group, list) and all(isinstance(item, int) for item in group):
+        group = _model.new_group(group)
+
+    if isinstance(_model, _SerialModel) or (group is not None and _model.get_rank() not in group):
+        return [tensor]
+
+    max_shape = torch.tensor(shapes).amax(dim=1)
+    padding_sizes = (max_shape - torch.tensor(tensor.shape)).tolist()
+    padded_tensor = torch.nn.functional.pad(
+        tensor, tuple(itertools.chain.from_iterable(map(lambda dim_size: (0, dim_size), reversed(padding_sizes))))
+    )
+    all_padded_tensors: torch.Tensor = _model.all_gather(padded_tensor, group=group)  # .split(max_shape[0], dim=0)
+    return [
+        all_padded_tensors[
+            [
+                slice(rank * max_shape[0] if dim == 0 else 0, rank * max_shape[0] + dim_size if dim == 0 else dim_size)
+                for dim, dim_size in enumerate(shape)
+            ]
+        ]
+        for rank, shape in enumerate(shapes)
+        if group is None or rank in group
+    ]
+
+
 def all_gather(
-    tensor: Union[torch.Tensor, float, str], group: Optional[Union[Any, List[int]]] = None
-) -> Union[torch.Tensor, float, List[float], List[str]]:
+    tensor: Union[torch.Tensor, float, str],
+    group: Optional[Union[Any, List[int]]] = None,
+    tensor_different_shape: bool = False,
+) -> Union[torch.Tensor, float, List[float], List[str], List[torch.Tensor]]:
     """Helper method to perform all gather operation.
 
     Args:
-        tensor: tensor or number or str to collect across participating processes.
+        tensor: tensor or number or str to collect across participating processes. If tensor, it should have
+            the same number of dimensions across processes.
         group: list of integer or the process group for each backend. If None, the default process group will be used.
+        tensor_different_shape: If True, it accounts for difference in input shape across processes. In this case, it
+            induces more collective operations. If False, `tensor` should have the same shape across processes.
+            Ignored when `tensor` is not a tensor. Default False.
+
 
     Returns:
-        torch.Tensor of shape ``(world_size * tensor.shape[0], tensor.shape[1], ...)`` if input is a tensor or
-        torch.Tensor of shape ``(world_size, )`` if input is a number or
-        List of strings if input is a string
+        If input is a tensor, returns a torch.Tensor of shape ``(world_size * tensor.shape[0], tensor.shape[1], ...)``
+        if ``tensor_different_shape = False``, otherwise a list of tensors with length ``world_size``(if ``group``
+        is `None`) or `len(group)`. If current process does not belong to `group`, a list with `tensor` as its only
+        item is retured.
+        If input is a number, a torch.Tensor of shape ``(world_size, )`` is returned and finally a list of strings
+        is returned if input is a string.
 
     .. versionchanged:: 0.4.11
         added ``group``
+
+    .. versionchanged:: 0.5.1
+        added ``tensor_different_shape``
     """
     if _need_to_sync and isinstance(_model, _SerialModel):
         sync(temporary=True)
 
     if isinstance(group, list) and all(isinstance(item, int) for item in group):
         group = _model.new_group(group)
 
+    if isinstance(tensor, torch.Tensor) and tensor_different_shape:
+        if isinstance(_model, _SerialModel) or (group is not None and _model.get_rank() not in group):
+            return [tensor]
+        all_shapes: torch.Tensor = _model.all_gather(torch.tensor(tensor.shape), group=group).view(
+            -1, len(tensor.shape)
+        )
+        return _all_gather_tensors_with_shapes(tensor, all_shapes.tolist(), group=group)
+
     return _model.all_gather(tensor, group=group)
 
 

ignite/metrics/__init__.py

@@ -9,6 +9,7 @@
 from ignite.metrics.gan.inception_score import InceptionScore
 from ignite.metrics.loss import Loss
 from ignite.metrics.mean_absolute_error import MeanAbsoluteError
+from ignite.metrics.mean_average_precision import MeanAveragePrecision
 from ignite.metrics.mean_pairwise_distance import MeanPairwiseDistance
 from ignite.metrics.mean_squared_error import MeanSquaredError
 from ignite.metrics.metric import BatchFiltered, BatchWise, EpochWise, Metric, MetricUsage
@@ -23,6 +24,7 @@
 from ignite.metrics.running_average import RunningAverage
 from ignite.metrics.ssim import SSIM
 from ignite.metrics.top_k_categorical_accuracy import TopKCategoricalAccuracy
+from ignite.metrics.vision.object_detection_map import ObjectDetectionMAP
 
 __all__ = [
     "Metric",
@@ -58,4 +60,6 @@
     "Rouge",
     "RougeN",
     "RougeL",
+    "MeanAveragePrecision",
+    "ObjectDetectionMAP",
 ]