Update median (#2681)

* Update median

* Update _base.py

* Update _base.py

* Update _base.py

* Update _base.py

* Update _base.py

* Update method name to `_torch_median`

* Update test__base.py

* Update test__base.py

* Update test__base.py

* Update test__base.py

* Update `get_rank` position

* Update test__base.py

* Update test_median_absolute_percentage_error.py

* Update test_median_relative_absolute_error.py

* Update test__base.py

Co-authored-by: vfdev <vfdev.5@gmail.com>

Author: puhuk

ignite/contrib/metrics/regression/_base.py

@@ -30,6 +30,16 @@ def _check_output_types(output: Tuple[torch.Tensor, torch.Tensor]) -> None:
         raise TypeError(f"Input y dtype should be float 16, 32 or 64, but given {y.dtype}")
 
 
+def _torch_median(output: torch.Tensor) -> float:
+    output = output.view(-1)
+    len_ = len(output)
+
+    if len_ % 2 == 0:
+        return float((torch.kthvalue(output, len_ // 2)[0] + torch.kthvalue(output, len_ // 2 + 1)[0]) / 2)
+    else:
+        return float(torch.kthvalue(output, len_ // 2 + 1)[0])
+
+
 class _BaseRegression(Metric):
     # Base class for all regression metrics
     # `update` method check the shapes and call internal overloaded

ignite/contrib/metrics/regression/median_absolute_error.py

@@ -2,12 +2,14 @@
 
 import torch
 
+from ignite.contrib.metrics.regression._base import _torch_median
+
 from ignite.metrics import EpochMetric
 
 
 def median_absolute_error_compute_fn(y_pred: torch.Tensor, y: torch.Tensor) -> float:
     e = torch.abs(y.view_as(y_pred) - y_pred)
-    return torch.median(e).item()
+    return _torch_median(e)
 
 
 class MedianAbsoluteError(EpochMetric):

ignite/contrib/metrics/regression/median_absolute_percentage_error.py

@@ -2,12 +2,14 @@
 
 import torch
 
+from ignite.contrib.metrics.regression._base import _torch_median
+
 from ignite.metrics import EpochMetric
 
 
 def median_absolute_percentage_error_compute_fn(y_pred: torch.Tensor, y: torch.Tensor) -> float:
     e = torch.abs(y.view_as(y_pred) - y_pred) / torch.abs(y.view_as(y_pred))
-    return 100.0 * torch.median(e).item()
+    return 100.0 * _torch_median(e)
 
 
 class MedianAbsolutePercentageError(EpochMetric):

ignite/contrib/metrics/regression/median_relative_absolute_error.py

@@ -2,12 +2,14 @@
 
 import torch
 
+from ignite.contrib.metrics.regression._base import _torch_median
+
 from ignite.metrics import EpochMetric
 
 
 def median_relative_absolute_error_compute_fn(y_pred: torch.Tensor, y: torch.Tensor) -> float:
     e = torch.abs(y.view_as(y_pred) - y_pred) / torch.abs(y.view_as(y_pred) - torch.mean(y))
-    return torch.median(e).item()
+    return _torch_median(e)
 
 
 class MedianRelativeAbsoluteError(EpochMetric):

tests/ignite/contrib/metrics/regression/test__base.py

@@ -1,7 +1,13 @@
+from typing import Optional
+
+import numpy as np
+
 import pytest
 import torch
 
-from ignite.contrib.metrics.regression._base import _BaseRegression
+import ignite.distributed as idist
+
+from ignite.contrib.metrics.regression._base import _BaseRegression, _torch_median
 
 
 def test_base_regression_shapes():
@@ -37,3 +43,57 @@ def compute(self):
     with pytest.raises(TypeError, match=r"Input y dtype should be float"):
         y = torch.tensor([1, 1])
         m.update((y.float(), y))
+
+
+@pytest.mark.parametrize("size", [100, 101, (30, 3), (31, 3)])
+def test_torch_median_numpy(size, device: Optional[str] = None):
+    data = torch.rand(size).to(device)
+    assert _torch_median(data) == np.median(data.cpu().numpy())
+
+
+@pytest.mark.parametrize("size", [101, (31, 3)])
+def test_torch_median_quantile(size, device: Optional[str] = None):
+    data = torch.rand(size).to(device)
+    assert _torch_median(data) == torch.quantile(data, 0.5, interpolation="midpoint")
+
+    size = 101
+    data = torch.rand(size=(size,))
+    assert _torch_median(data) == torch.median(data)
+
+
+@pytest.mark.tpu
+@pytest.mark.parametrize("size", [100, 101, (30, 3), (31, 3)])
+@pytest.mark.skipif(not idist.has_xla_support, reason="Skip if no PyTorch XLA package")
+def test_on_even_size_xla(size):
+    device = "xla"
+    test_torch_median_numpy(size, device=device)
+
+
+@pytest.mark.parametrize("size", [100, 101, (30, 3), (31, 3)])
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="Skip if no GPU")
+def test_on_even_size_gpu(size):
+    test_torch_median_numpy(size, device="cuda")
+
+
+@pytest.mark.parametrize("size", [100, 101, (30, 3), (31, 3)])
+def test_create_even_size_cpu(size):
+    test_torch_median_numpy(size, device="cpu")
+
+
+@pytest.mark.tpu
+@pytest.mark.parametrize("size", [101, (31, 3)])
+@pytest.mark.skipif(not idist.has_xla_support, reason="Skip if no PyTorch XLA package")
+def test_on_odd_size_xla(size):
+    device = "xla"
+    test_torch_median_quantile(size, device=device)
+
+
+@pytest.mark.parametrize("size", [101, (31, 3)])
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="Skip if no GPU")
+def test_on_odd_size_gpu(size):
+    test_torch_median_quantile(size, device="cuda")
+
+
+@pytest.mark.parametrize("size", [101, (31, 3)])
+def test_create_odd_size_cpu(size):
+    test_torch_median_quantile(size, device="cpu")

tests/ignite/contrib/metrics/regression/test_median_absolute_error.py

@@ -108,12 +108,9 @@ def update_fn(engine, batch):
 
 
 def _test_distrib_compute(device):
-    rank = idist.get_rank()
-
     def _test(metric_device):
         metric_device = torch.device(metric_device)
         m = MedianAbsoluteError(device=metric_device)
-        torch.manual_seed(10 + rank)
 
         size = 105
 
@@ -132,39 +129,29 @@ def _test(metric_device):
 
         e = np.abs(np_y - np_y_pred)
 
-        # The results between numpy.median() and torch.median() are Inconsistant
-        # when the length of the array/tensor is even. So this is a hack to avoid that.
-        # issue: https://github.com/pytorch/pytorch/issues/1837
-        if np_y_pred.shape[0] % 2 == 0:
-            e_prepend = np.insert(e, 0, e[0], axis=0)
-            np_res_prepend = np.median(e_prepend)
-            assert pytest.approx(res) == np_res_prepend
-        else:
-            np_res = np.median(e)
-            assert pytest.approx(res) == np_res
-
-    for _ in range(3):
+        np_res = np.median(e)
+        assert pytest.approx(res) == np_res
+
+    rank = idist.get_rank()
+    for i in range(3):
+        torch.manual_seed(10 + rank + i)
         _test("cpu")
         if device.type != "xla":
             _test(idist.device())
 
 
 def _test_distrib_integration(device):
-
-    rank = idist.get_rank()
-    torch.manual_seed(12)
-
     def _test(n_epochs, metric_device):
         metric_device = torch.device(metric_device)
         n_iters = 80
-        size = 105
-        y_true = torch.rand(size=(size,)).to(device)
-        y_preds = torch.rand(size=(size,)).to(device)
+        batch_size = 105
+        y_true = torch.rand(size=(n_iters * batch_size,)).to(device)
+        y_preds = torch.rand(size=(n_iters * batch_size,)).to(device)
 
         def update(engine, i):
             return (
-                y_preds[i * size : (i + 1) * size],
-                y_true[i * size : (i + 1) * size],
+                y_preds[i * batch_size : (i + 1) * batch_size],
+                y_true[i * batch_size : (i + 1) * batch_size],
             )
 
         engine = Engine(update)
@@ -175,6 +162,9 @@ def update(engine, i):
         data = list(range(n_iters))
         engine.run(data=data, max_epochs=n_epochs)
 
+        y_preds = idist.all_gather(y_preds)
+        y_true = idist.all_gather(y_true)
+
         assert "mae" in engine.state.metrics
 
         res = engine.state.metrics["mae"]
@@ -191,7 +181,9 @@ def update(engine, i):
     if device.type != "xla":
         metric_devices.append(idist.device())
     for metric_device in metric_devices:
-        for _ in range(2):
+        rank = idist.get_rank()
+        for i in range(2):
+            torch.manual_seed(10 + rank + i)
             _test(n_epochs=1, metric_device=metric_device)
             _test(n_epochs=2, metric_device=metric_device)
 

tests/ignite/contrib/metrics/regression/test_median_absolute_percentage_error.py

@@ -108,12 +108,9 @@ def update_fn(engine, batch):
 
 
 def _test_distrib_compute(device):
-    rank = idist.get_rank()
-
     def _test(metric_device):
         metric_device = torch.device(metric_device)
         m = MedianAbsolutePercentageError(device=metric_device)
-        torch.manual_seed(10 + rank)
 
         size = 105
 
@@ -133,34 +130,24 @@ def _test(metric_device):
 
         e = np.abs(np_y - np_y_pred) / np.abs(np_y)
 
-        # The results between numpy.median() and torch.median() are Inconsistant
-        # when the length of the array/tensor is even. So this is a hack to avoid that.
-        # issue: https://github.com/pytorch/pytorch/issues/1837
-        if np_y_pred.shape[0] % 2 == 0:
-            e_prepend = np.insert(e, 0, e[0], axis=0)
-            np_res_prepend = 100.0 * np.median(e_prepend)
-            assert pytest.approx(res) == np_res_prepend
-        else:
-            np_res = 100.0 * np.median(e)
-            assert pytest.approx(res) == np_res
-
-    for _ in range(3):
+        np_res = 100.0 * np.median(e)
+        assert pytest.approx(res) == np_res
+
+    rank = idist.get_rank()
+    for i in range(3):
+        torch.manual_seed(10 + rank + i)
         _test("cpu")
         if device.type != "xla":
             _test(idist.device())
 
 
 def _test_distrib_integration(device):
-
-    rank = idist.get_rank()
-    torch.manual_seed(12)
-
     def _test(n_epochs, metric_device):
         metric_device = torch.device(metric_device)
         n_iters = 80
         size = 105
-        y_true = torch.rand(size=(size,)).to(device)
-        y_preds = torch.rand(size=(size,)).to(device)
+        y_true = torch.rand(size=(n_iters * size,)).to(device)
+        y_preds = torch.rand(size=(n_iters * size,)).to(device)
 
         def update(engine, i):
             return (
@@ -176,6 +163,9 @@ def update(engine, i):
         data = list(range(n_iters))
         engine.run(data=data, max_epochs=n_epochs)
 
+        y_preds = idist.all_gather(y_preds)
+        y_true = idist.all_gather(y_true)
+
         assert "mape" in engine.state.metrics
 
         res = engine.state.metrics["mape"]
@@ -186,22 +176,15 @@ def update(engine, i):
         e = np.abs(np_y_true - np_y_preds) / np.abs(np_y_true)
         np_res = 100.0 * np.median(e)
 
-        e_prepend = np.insert(e, 0, e[0], axis=0)
-        np_res_prepend = 100.0 * np.median(e_prepend)
-
-        # The results between numpy.median() and torch.median() are Inconsistant
-        # when the length of the array/tensor is even. So this is a hack to avoid that.
-        # issue: https://github.com/pytorch/pytorch/issues/1837
-        if np_y_preds.shape[0] % 2 == 0:
-            assert pytest.approx(res) == np_res_prepend
-        else:
-            assert pytest.approx(res) == np_res
+        assert pytest.approx(res) == np_res
 
     metric_devices = ["cpu"]
     if device.type != "xla":
         metric_devices.append(idist.device())
     for metric_device in metric_devices:
-        for _ in range(2):
+        rank = idist.get_rank()
+        for i in range(2):
+            torch.manual_seed(12 + rank + i)
             _test(n_epochs=1, metric_device=metric_device)
             _test(n_epochs=2, metric_device=metric_device)
 

tests/ignite/contrib/metrics/regression/test_median_relative_absolute_error.py

@@ -108,8 +108,6 @@ def update_fn(engine, batch):
 
 
 def _test_distrib_compute(device):
-    rank = idist.get_rank()
-
     def _test(metric_device):
         metric_device = torch.device(metric_device)
         m = MedianRelativeAbsoluteError(device=metric_device)
@@ -133,28 +131,17 @@ def _test(metric_device):
 
         e = np.abs(np_y - np_y_pred) / np.abs(np_y - np_y.mean())
 
-        # The results between numpy.median() and torch.median() are Inconsistant
-        # when the length of the array/tensor is even. So this is a hack to avoid that.
-        # issue: https://github.com/pytorch/pytorch/issues/1837
-        if np_y_pred.shape[0] % 2 == 0:
-            e_prepend = np.insert(e, 0, e[0], axis=0)
-            np_res_prepend = np.median(e_prepend)
-            assert pytest.approx(res) == np_res_prepend
-        else:
-            np_res = np.median(e)
-            assert pytest.approx(res) == np_res
+        np_res = np.median(e)
+        assert pytest.approx(res) == np_res
 
+    rank = idist.get_rank()
     for _ in range(3):
         _test("cpu")
         if device.type != "xla":
             _test(idist.device())
 
 
 def _test_distrib_integration(device):
-
-    rank = idist.get_rank()
-    torch.manual_seed(12)
-
     def _test(n_epochs, metric_device):
         metric_device = torch.device(metric_device)
         n_iters = 80
@@ -176,6 +163,9 @@ def update(engine, i):
         data = list(range(n_iters))
         engine.run(data=data, max_epochs=n_epochs)
 
+        y_true = idist.all_gather(y_true)
+        y_preds = idist.all_gather(y_preds)
+
         assert "mare" in engine.state.metrics
 
         res = engine.state.metrics["mare"]
@@ -192,7 +182,9 @@ def update(engine, i):
     if device.type != "xla":
         metric_devices.append(idist.device())
     for metric_device in metric_devices:
-        for _ in range(2):
+        rank = idist.get_rank()
+        for i in range(2):
+            torch.manual_seed(12 + rank + i)
             _test(n_epochs=1, metric_device=metric_device)
             _test(n_epochs=2, metric_device=metric_device)