Fixes for pytorch<2.0 in average precision

Author: vfdev-5

ignite/metrics/mean_average_precision.py

@@ -2,6 +2,7 @@
 from typing import Callable, cast, List, Optional, Sequence, Tuple, Union
 
 import torch
+from packaging.version import Version
 from typing_extensions import Literal
 
 import ignite.distributed as idist
@@ -11,6 +12,9 @@
 from ignite.utils import to_onehot
 
 
+_torch_version_lt_113 = Version(torch.__version__) < Version("1.13.0")
+
+
 class _BaseAveragePrecision:
     def __init__(
         self,
@@ -97,9 +101,12 @@ def _compute_average_precision(self, recall: torch.Tensor, precision: torch.Tens
         if self.rec_thresholds is not None:
             rec_thresholds = self.rec_thresholds.repeat((*recall.shape[:-1], 1))
             rec_thresh_indices = torch.searchsorted(recall, rec_thresholds)
-            precision = precision.take_along_dim(
-                rec_thresh_indices.where(rec_thresh_indices != recall.size(-1), 0), dim=-1
-            ).where(rec_thresh_indices != recall.size(-1), 0)
+            rec_mask = rec_thresh_indices != recall.size(-1)
+            precision = torch.where(
+                rec_mask,
+                precision.take_along_dim(torch.where(rec_mask, rec_thresh_indices, 0), dim=-1),
+                0.0,
+            )
             recall = rec_thresholds
         recall_differential = recall.diff(
             dim=-1, prepend=torch.zeros((*recall.shape[:-1], 1), device=recall.device, dtype=recall.dtype)
@@ -335,9 +342,10 @@ def _compute_recall_and_precision(
         Returns:
             `(recall, precision)`
         """
-        indices = torch.argsort(y_pred, stable=True, descending=True)
+        kwargs = {} if _torch_version_lt_113 else {"stable": True}
+        indices = torch.argsort(y_pred, descending=True, **kwargs)
         tp_summation = y_true[indices].cumsum(dim=0)
-        if tp_summation.device != torch.device("mps"):
+        if tp_summation.device.type != "mps":
             tp_summation = tp_summation.double()
 
         # Adopted from Scikit-learn's implementation
@@ -371,7 +379,7 @@ def compute(self) -> Union[torch.Tensor, float]:
             torch.long if self._type == "multiclass" else torch.uint8,
             self._device,
         )
-        fp_precision = torch.double if self._device != torch.device("mps") else torch.float32
+        fp_precision = torch.double if self._device.type != "mps" else torch.float32
         y_pred = _cat_and_agg_tensors(self._y_pred, (num_classes,), fp_precision, self._device)
 
         if self._type == "multiclass":

ignite/metrics/vision/object_detection_average_precision_recall.py

@@ -1,6 +1,7 @@
 from typing import Callable, cast, Dict, List, Optional, Sequence, Tuple, Union
 
 import torch
+from packaging.version import Version
 from typing_extensions import Literal
 
 from ignite.metrics import MetricGroup
@@ -9,6 +10,9 @@
 from ignite.metrics.metric import Metric, reinit__is_reduced, sync_all_reduce
 
 
+_torch_version_lt_113 = Version(torch.__version__) < Version("1.13.0")
+
+
 def coco_tensor_list_to_dict_list(
     output: Tuple[
         Union[List[torch.Tensor], List[Dict[str, torch.Tensor]]],
@@ -213,7 +217,8 @@ def _compute_recall_and_precision(
         Returns:
             `(recall, precision)`
         """
-        indices = torch.argsort(scores, dim=-1, stable=True, descending=True)
+        kwargs = {} if _torch_version_lt_113 else {"stable": True}
+        indices = torch.argsort(scores, descending=True, **kwargs)
         tp = TP[..., indices]
         tp_summation = tp.cumsum(dim=-1)
         if tp_summation.device.type != "mps":
@@ -226,7 +231,7 @@ def _compute_recall_and_precision(
 
         recall = tp_summation / y_true_count
         predicted_positive = tp_summation + fp_summation
-        precision = tp_summation / torch.where(predicted_positive == 0, 1, predicted_positive)
+        precision = tp_summation / torch.where(predicted_positive == 0, 1.0, predicted_positive)
 
         return recall, precision
 
@@ -258,9 +263,12 @@ def _compute_average_precision(self, recall: torch.Tensor, precision: torch.Tens
             if recall.size(-1) != 0
             else torch.LongTensor([], device=self._device)
         )
-        precision_integrand = precision_integrand.take_along_dim(
-            rec_thresh_indices.where(rec_thresh_indices != recall.size(-1), 0), dim=-1
-        ).where(rec_thresh_indices != recall.size(-1), 0)
+        recall_mask = rec_thresh_indices != recall.size(-1)
+        precision_integrand = torch.where(
+            recall_mask,
+            precision_integrand.take_along_dim(torch.where(recall_mask, rec_thresh_indices, 0), dim=-1),
+            0.0,
+        )
         return torch.sum(precision_integrand, dim=-1) / len(cast(torch.Tensor, self.rec_thresholds))
 
     @reinit__is_reduced
@@ -298,6 +306,7 @@ def update(self, output: Tuple[List[Dict[str, torch.Tensor]], List[Dict[str, tor
                                             This key is optional.
                 ========= ================= =================================================
         """
+        kwargs = {} if _torch_version_lt_113 else {"stable": True}
         self._check_matching_input(output)
         for pred, target in zip(*output):
             labels = target["labels"]
@@ -312,7 +321,7 @@ def update(self, output: Tuple[List[Dict[str, torch.Tensor]], List[Dict[str, tor
 
             # Matching logic of object detection mAP, according to COCO reference implementation.
             if len(pred["labels"]):
-                best_detections_index = torch.argsort(pred["scores"], stable=True, descending=True)
+                best_detections_index = torch.argsort(pred["scores"], descending=True, **kwargs)
                 max_best_detections_index = torch.cat(
                     [
                         best_detections_index[pred["labels"][best_detections_index] == c][

tests/ignite/metrics/vision/test_object_detection_map.py

@@ -864,7 +864,7 @@ def test__compute_recall_and_precision():
 def test_compute(sample):
     device = idist.device()
 
-    if device == torch.device("mps"):
+    if device.type == "mps":
         pytest.skip("Due to MPS backend out of memory")
 
     # [email protected], [email protected], [email protected], AP-S, AP-M, AP-L, AR-1, AR-10, AR-100, AR-S, AR-M, AR-L
@@ -924,7 +924,7 @@ def test_integration(sample):
     bs = 3
 
     device = idist.device()
-    if device == torch.device("mps"):
+    if device.type == "mps":
         pytest.skip("Due to MPS backend out of memory")
 
     def update(engine, i):
@@ -995,7 +995,7 @@ def test_distrib_update_compute(distributed, sample):
 
     device = idist.device()
 
-    if device == torch.device("mps"):
+    if device.type == "mps":
         pytest.skip("Due to MPS backend out of memory")
 
     metric_device = "cpu" if device.type == "xla" else device