Update pairwise_dist, mse and lambda in test for generating data with different rank

tests/ignite/metrics/test_mean_pairwise_distance.py

@@ -56,22 +56,22 @@ def _test_distrib_integration(device):
     from ignite.engine import Engine
 
     rank = idist.get_rank()
-    torch.manual_seed(12)
+    torch.manual_seed(12 + rank)
 
-    n_iters = 100
-    s = 50
-    offset = n_iters * s
+    def _test(metric_device):
 
-    y_true = torch.rand(offset * idist.get_world_size(), 10).to(device)
-    y_preds = torch.rand(offset * idist.get_world_size(), 10).to(device)
+        n_iters = 100
+        batch_size = 50
 
-    def update(engine, i):
-        return (
-            y_preds[i * s + offset * rank : (i + 1) * s + offset * rank, ...],
-            y_true[i * s + offset * rank : (i + 1) * s + offset * rank, ...],
-        )
+        y_true = torch.rand(n_iters * batch_size, 10).to(device)
+        y_preds = torch.rand(n_iters * batch_size, 10).to(device)
+
+        def update(engine, i):
+            return (
+                y_preds[i * batch_size : (i + 1) * batch_size, ...],
+                y_true[i * batch_size : (i + 1) * batch_size, ...],
+            )
 
-    def _test(metric_device):
         engine = Engine(update)
 
         m = MeanPairwiseDistance(device=metric_device)
@@ -80,14 +80,20 @@ def _test(metric_device):
         data = list(range(n_iters))
         engine.run(data=data, max_epochs=1)
 
+        y_preds = idist.all_gather(y_preds)
+        y_true = idist.all_gather(y_true)
+
         assert "mpwd" in engine.state.metrics
         res = engine.state.metrics["mpwd"]
 
         true_res = []
         for i in range(n_iters * idist.get_world_size()):
             true_res.append(
                 torch.pairwise_distance(
-                    y_true[i * s : (i + 1) * s, ...], y_preds[i * s : (i + 1) * s, ...], p=m._p, eps=m._eps
+                    y_true[i * batch_size : (i + 1) * batch_size, ...],
+                    y_preds[i * batch_size : (i + 1) * batch_size, ...],
+                    p=m._p,
+                    eps=m._eps,
                 )
                 .cpu()
                 .numpy()

tests/ignite/metrics/test_mean_squared_error.py

@@ -58,20 +58,21 @@ def _test_distrib_integration(device, tol=1e-6):
     from ignite.engine import Engine
 
     rank = idist.get_rank()
-    n_iters = 100
-    s = 10
-    offset = n_iters * s
+    torch.manual_seed(12 + rank)
 
-    y_true = torch.arange(0, offset * idist.get_world_size(), dtype=torch.float).to(device)
-    y_preds = torch.ones(offset * idist.get_world_size(), dtype=torch.float).to(device)
+    def _test(metric_device):
+        n_iters = 100
+        batch_size = 10
 
-    def update(engine, i):
-        return (
-            y_preds[i * s + offset * rank : (i + 1) * s + offset * rank],
-            y_true[i * s + offset * rank : (i + 1) * s + offset * rank],
-        )
+        y_true = torch.arange(0, n_iters * batch_size, dtype=torch.float).to(device)
+        y_preds = torch.ones(n_iters * batch_size, dtype=torch.float).to(device)
+
+        def update(engine, i):
+            return (
+                y_preds[i * batch_size : (i + 1) * batch_size],
+                y_true[i * batch_size : (i + 1) * batch_size],
+            )
 
-    def _test(metric_device):
         engine = Engine(update)
 
         m = MeanSquaredError(device=metric_device)
@@ -80,6 +81,9 @@ def _test(metric_device):
         data = list(range(n_iters))
         engine.run(data=data, max_epochs=1)
 
+        y_preds = idist.all_gather(y_preds)
+        y_true = idist.all_gather(y_true)
+
         assert "mse" in engine.state.metrics
         res = engine.state.metrics["mse"]
 

tests/ignite/metrics/test_metrics_lambda.py

@@ -437,29 +437,25 @@ def compute_true_somemetric(y_pred, y):
 def _test_distrib_integration(device):
 
     rank = idist.get_rank()
-    np.random.seed(12)
 
     n_iters = 10
     batch_size = 10
     n_classes = 10
 
     def _test(metric_device):
-        y_true = np.arange(0, n_iters * batch_size * idist.get_world_size(), dtype="int64") % n_classes
-        y_pred = 0.2 * np.random.rand(n_iters * batch_size * idist.get_world_size(), n_classes)
-        for i in range(n_iters * batch_size * idist.get_world_size()):
+        y_true = torch.arange(0, n_iters * batch_size, dtype=torch.int64).to(device) % n_classes
+        y_pred = 0.2 * torch.rand(n_iters * batch_size, n_classes).to(device)
+        for i in range(n_iters * batch_size):
             if np.random.rand() > 0.4:
                 y_pred[i, y_true[i]] = 1.0
             else:
                 j = np.random.randint(0, n_classes)
                 y_pred[i, j] = 0.7
 
-        y_true = y_true.reshape(n_iters * idist.get_world_size(), batch_size)
-        y_pred = y_pred.reshape(n_iters * idist.get_world_size(), batch_size, n_classes)
-
         def update_fn(engine, i):
-            y_true_batch = y_true[i + rank * n_iters, ...]
-            y_pred_batch = y_pred[i + rank * n_iters, ...]
-            return torch.from_numpy(y_pred_batch), torch.from_numpy(y_true_batch)
+            y_true_batch = y_true[i * batch_size : (i + 1) * batch_size, ...]
+            y_pred_batch = y_pred[i * batch_size : (i + 1) * batch_size, ...]
+            return y_pred_batch, y_true_batch
 
         evaluator = Engine(update_fn)
 
@@ -478,13 +474,17 @@ def Fbeta(r, p, beta):
         data = list(range(n_iters))
         state = evaluator.run(data, max_epochs=1)
 
+        y_pred = idist.all_gather(y_pred)
+        y_true = idist.all_gather(y_true)
+
         assert "f1" in state.metrics
         assert "ff1" in state.metrics
         f1_true = f1_score(y_true.ravel(), np.argmax(y_pred.reshape(-1, n_classes), axis=-1), average="macro")
         assert f1_true == approx(state.metrics["f1"])
         assert 1.0 + f1_true == approx(state.metrics["ff1"])
 
-    for _ in range(3):
+    for i in range(3):
+        torch.manual_seed(12 + rank + i)
         _test("cpu")
         if device.type != "xla":
             _test(idist.device())
@@ -493,17 +493,17 @@ def Fbeta(r, p, beta):
 def _test_distrib_metrics_on_diff_devices(device):
     n_classes = 10
     n_iters = 12
-    s = 16
-    offset = n_iters * s
+    batch_size = 16
     rank = idist.get_rank()
+    torch.manual_seed(12 + rank)
 
-    y_true = torch.randint(0, n_classes, size=(offset * idist.get_world_size(),)).to(device)
-    y_preds = torch.rand(offset * idist.get_world_size(), n_classes).to(device)
+    y_true = torch.randint(0, n_classes, size=(n_iters * batch_size,)).to(device)
+    y_preds = torch.rand(n_iters * batch_size, n_classes).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],
         )
 
     precision = Precision(average=False, device="cpu")