Add utils mpi_scatter_from_rank0 (#1298)

Author: pescap

deepxde/data/pde.py

@@ -4,7 +4,7 @@
 from .. import backend as bkd
 from .. import config
 from ..backend import backend_name
-from ..utils import get_num_args, run_if_all_none
+from ..utils import get_num_args, run_if_all_none, mpi_scatter_from_rank0
 
 
 class PDE(Data):
@@ -186,22 +186,7 @@ def train_next_batch(self, batch_size=None):
             config.comm.Bcast(self.train_x_bc, root=0)
         self.train_x = self.train_x_bc
         if config.parallel_scaling == "strong":
-            # Split the training points over each rank.
-            # We drop last points in order to have the same number of points per rank
-            if len(self.train_x_all) < config.world_size:
-                raise ValueError(
-                    "The number of training points is smaller than the number of processes. Please use more points."
-                )
-            train_x_all_shape = list(
-                self.train_x_all.shape
-            )  # We transform to list to support item assignment
-            num_split = train_x_all_shape[0] // config.world_size
-            train_x_all_shape[0] = num_split
-            train_x_all_split = np.empty(
-                train_x_all_shape, dtype=self.train_x_all.dtype
-            )
-            config.comm.Scatter(self.train_x_all, train_x_all_split, root=0)
-            self.train_x_all = train_x_all_split
+            self.train_x_all = mpi_scatter_from_rank0(self.train_x_all)
         if self.pde is not None:
             self.train_x = np.vstack((self.train_x, self.train_x_all))
         self.train_y = self.soln(self.train_x) if self.soln else None

deepxde/utils/internal.py

@@ -200,3 +200,34 @@ def get_num_args(func):
     # g = dummy(a.f)
     params = inspect.signature(func).parameters
     return len(params) - ("self" in params)
+
+
+def mpi_scatter_from_rank0(array, drop_last=True):
+    """Scatter the given array into continuous subarrays of equal size from rank 0 to all ranks.
+
+    Args:
+        array: Numpy array to be split.
+        drop_last (bool): Whether to discard the remainder samples
+            not divisible by world_size. Default: True.
+
+    Returns:
+        array: Scattered Numpy array.
+    """
+    # TODO: support drop_last=False
+    if config.world_size == 1:
+        return array
+    if not drop_last:
+        raise ValueError("Only support drop_last=True now.")
+    if len(array) < config.world_size:
+        raise ValueError(
+            "The number of training points is smaller than the number of processes. Please use more points."
+        )
+    array_shape = list(array.shape)  # We transform to list to support item assignment
+    num_split = array_shape[0] // config.world_size
+    array_shape[0] = num_split
+    array_split = np.empty(array_shape, dtype=array.dtype)
+    array = array[
+        : num_split * config.world_size
+    ]  # We truncate array size to be a multiple of num_split to prevent a MPI error.
+    config.comm.Scatter(array, array_split, root=0)
+    return array_split