Update geometric_mean_relative_absolute_error, manhattan_distance…

…, `maximum_absolute_error` (#2664)

Update `geometric_mean_relative_absolute_error`, `manhattan_distance`, `maximum_absolute_error` to generate with different rank

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

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

@@ -79,12 +79,10 @@ def update_fn(engine, batch):
 def _test_distrib_compute(device):
 
     rank = idist.get_rank()
-    torch.manual_seed(12)
 
     def _test(metric_device):
         metric_device = torch.device(metric_device)
         m = GeometricMeanRelativeAbsoluteError(device=metric_device)
-        torch.manual_seed(10 + rank)
 
         y_pred = torch.rand(size=(100,), device=device)
         y = torch.rand(size=(100,), device=device)
@@ -101,7 +99,8 @@ def _test(metric_device):
 
         assert m.compute() == pytest.approx(np_gmrae, rel=1e-4)
 
-    for _ in range(3):
+    for i in range(3):
+        torch.manual_seed(12 + rank + i)
         _test("cpu")
         if device.type != "xla":
             _test(idist.device())
@@ -115,17 +114,15 @@ def _test_distrib_integration(device):
     def _test(n_epochs, metric_device):
         metric_device = torch.device(metric_device)
         n_iters = 80
-        s = 16
-        n_classes = 2
+        batch_size = 16
 
-        offset = n_iters * s
-        y_true = torch.rand(size=(offset * idist.get_world_size(),)).to(device)
-        y_preds = torch.rand(size=(offset * idist.get_world_size(),)).to(device)
+        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 * s + rank * offset : (i + 1) * s + rank * offset],
-                y_true[i * s + rank * offset : (i + 1) * s + rank * offset],
+                y_preds[i * batch_size : (i + 1) * batch_size],
+                y_true[i * batch_size : (i + 1) * batch_size],
             )
 
         engine = Engine(update)
@@ -136,6 +133,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 "gmrae" in engine.state.metrics
 
         res = engine.state.metrics["gmrae"]
@@ -151,7 +151,8 @@ def update(engine, i):
     if device.type != "xla":
         metric_devices.append(idist.device())
     for metric_device in metric_devices:
-        for _ in range(2):
+        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_manhattan_distance.py

@@ -109,7 +109,6 @@ def _test_distrib_compute(device):
     def _test(metric_device):
         metric_device = torch.device(metric_device)
         m = ManhattanDistance(device=metric_device)
-        torch.manual_seed(10 + rank)
 
         y_pred = torch.randint(0, 10, size=(10,), device=device).float()
         y = torch.randint(0, 10, size=(10,), device=device).float()
@@ -125,7 +124,8 @@ def _test(metric_device):
         res = m.compute()
         assert manhattan.pairwise([np_y_pred, np_y])[0][1] == pytest.approx(res)
 
-    for _ in range(3):
+    for i in range(3):
+        torch.manual_seed(10 + rank + i)
         _test("cpu")
         if device.type != "xla":
             _test(idist.device())
@@ -134,24 +134,21 @@ def _test(metric_device):
 def _test_distrib_integration(device):
 
     rank = idist.get_rank()
-    torch.manual_seed(12)
 
     manhattan = DistanceMetric.get_metric("manhattan")
 
     def _test(n_epochs, metric_device):
         metric_device = torch.device(metric_device)
         n_iters = 80
-        s = 16
-        n_classes = 2
+        batch_size = 16
 
-        offset = n_iters * s
-        y_true = torch.rand(size=(offset * idist.get_world_size(),)).to(device)
-        y_preds = torch.rand(size=(offset * idist.get_world_size(),)).to(device)
+        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 * s + rank * offset : (i + 1) * s + rank * offset],
-                y_true[i * s + rank * offset : (i + 1) * s + rank * offset],
+                y_preds[i * batch_size : (i + 1) * batch_size],
+                y_true[i * batch_size : (i + 1) * batch_size],
             )
 
         engine = Engine(update)
@@ -162,6 +159,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 "md" in engine.state.metrics
 
         res = engine.state.metrics["md"]
@@ -177,7 +177,8 @@ def update(engine, i):
     if device.type != "xla":
         metric_devices.append(idist.device())
     for metric_device in metric_devices:
-        for _ in range(2):
+        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_maximum_absolute_error.py

@@ -103,7 +103,6 @@ def _test_distrib_compute(device):
     def _test(metric_device):
         metric_device = torch.device(metric_device)
         m = MaximumAbsoluteError(device=metric_device)
-        torch.manual_seed(10 + rank)
 
         y_pred = torch.randint(0, 10, size=(10,), device=device).float()
         y = torch.randint(0, 10, size=(10,), device=device).float()
@@ -123,7 +122,8 @@ def _test(metric_device):
 
         assert np_max == pytest.approx(res)
 
-    for _ in range(3):
+    for i in range(3):
+        torch.manual_seed(10 + rank + i)
         _test("cpu")
         if device.type != "xla":
             _test(idist.device())
@@ -132,22 +132,19 @@ def _test(metric_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
-        s = 16
-        n_classes = 2
+        batch_size = 16
 
-        offset = n_iters * s
-        y_true = torch.rand(size=(offset * idist.get_world_size(),)).to(device)
-        y_preds = torch.rand(size=(offset * idist.get_world_size(),)).to(device)
+        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 * s + rank * offset : (i + 1) * s + rank * offset],
-                y_true[i * s + rank * offset : (i + 1) * s + rank * offset],
+                y_preds[i * batch_size : (i + 1) * batch_size],
+                y_true[i * batch_size : (i + 1) * batch_size],
             )
 
         engine = Engine(update)
@@ -158,6 +155,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"]
@@ -173,7 +173,8 @@ def update(engine, i):
     if device.type != "xla":
         metric_devices.append(idist.device())
     for metric_device in metric_devices:
-        for _ in range(2):
+        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)