Update rmse, ssim, top_k_categorical_accuracy in test for generating data with different rank

tests/ignite/metrics/test_root_mean_squared_error.py

@@ -58,17 +58,19 @@ 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
 
-    y_true = torch.arange(0, offset * idist.get_world_size(), dtype=torch.float).to(device)
-    y_preds = (rank + 1) * torch.ones(offset, dtype=torch.float).to(device)
+    def _test(metric_device):
+        n_iters = 2
+        batch_size = 3
 
-    def update(engine, i):
-        return y_preds[i * s : (i + 1) * s], y_true[i * s + offset * rank : (i + 1) * s + offset * rank]
+        torch.manual_seed(12 + rank)
+
+        y_true = torch.arange(0, n_iters * batch_size, dtype=torch.float).to(device)
+        y_preds = (rank + 1) * 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 = RootMeanSquaredError(device=metric_device)
@@ -77,15 +79,13 @@ 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 "rmse" in engine.state.metrics
         res = engine.state.metrics["rmse"]
 
-        y_preds_full = []
-        for i in range(idist.get_world_size()):
-            y_preds_full.append((i + 1) * torch.ones(offset))
-        y_preds_full = torch.stack(y_preds_full).to(device).flatten()
-
-        true_res = np.sqrt(np.mean(np.square((y_true - y_preds_full).cpu().numpy())))
+        true_res = np.sqrt(np.mean(np.square((y_true - y_preds).cpu().numpy())))
 
         assert pytest.approx(res, rel=tol) == true_res
 

tests/ignite/metrics/test_ssim.py

@@ -130,18 +130,18 @@ def _test_distrib_integration(device, tol=1e-4):
     from ignite.engine import Engine
 
     rank = idist.get_rank()
+    torch.manual_seed(12 + rank)
     n_iters = 100
-    s = 10
-    offset = n_iters * s
+    batch_size = 10
 
     def _test(metric_device):
-        y_pred = torch.rand(offset * idist.get_world_size(), 3, 28, 28, dtype=torch.float, device=device)
+        y_pred = torch.rand(n_iters * batch_size, 3, 28, 28, dtype=torch.float, device=device)
         y = y_pred * 0.65
 
         def update(engine, i):
             return (
-                y_pred[i * s + offset * rank : (i + 1) * s + offset * rank],
-                y[i * s + offset * rank : (i + 1) * s + offset * rank],
+                y_pred[i * batch_size : (i + 1) * batch_size],
+                y[i * batch_size : (i + 1) * batch_size],
             )
 
         engine = Engine(update)
@@ -150,6 +150,9 @@ def update(engine, i):
         data = list(range(n_iters))
         engine.run(data=data, max_epochs=1)
 
+        y_pred = idist.all_gather(y_pred)
+        y = idist.all_gather(y)
+
         assert "ssim" in engine.state.metrics
         res = engine.state.metrics["ssim"]
 

tests/ignite/metrics/test_top_k_categorical_accuracy.py

@@ -58,22 +58,18 @@ def top_k_accuracy(y_true, y_pred, k=5, normalize=True):
 def _test_distrib_integration(device):
     from ignite.engine import Engine
 
-    rank = idist.get_rank()
-    torch.manual_seed(12)
-
     def _test(n_epochs, metric_device):
         n_iters = 100
-        s = 16
+        batch_size = 16
         n_classes = 10
 
-        offset = n_iters * s
-        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],
             )
 
         engine = Engine(update)
@@ -85,6 +81,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 "acc" in engine.state.metrics
         res = engine.state.metrics["acc"]
         if isinstance(res, torch.Tensor):
@@ -97,7 +96,9 @@ def update(engine, i):
     metric_devices = ["cpu"]
     if device.type != "xla":
         metric_devices.append(idist.device())
-    for _ in range(3):
+    rank = idist.get_rank()
+    for i in range(3):
+        torch.manual_seed(12 + rank + i)
         for metric_device in metric_devices:
             _test(n_epochs=1, metric_device=metric_device)
             _test(n_epochs=2, metric_device=metric_device)