Add parallel scaling parameter dde.config.parallel_scaling (#1279)

Author: pescap

deepxde/config.py

@@ -7,6 +7,8 @@
 from .backend import backend_name, tf, torch, paddle
 from .real import Real
 
+# Data parallel
+parallel_scaling = None
 # Data parallel via Horovod
 hvd = None
 comm = None
@@ -21,9 +23,10 @@
         if world_size > 1:
             from mpi4py import MPI
 
+            parallel_scaling = "weak"
             comm = MPI.COMM_WORLD
             tf.compat.v1.disable_eager_execution()  # Without this line, Horovod broadcasting fails.
-            rank = hvd.rank() # Only single node acceleration supported so far.
+            rank = hvd.rank()  # Only single node acceleration supported so far.
             if rank == 0:
                 print(f"\nParallel training with {world_size} processes.\n")
         else:
@@ -198,3 +201,15 @@ def disable_xla_jit():
     This is equivalent with ``enable_xla_jit(False)``.
     """
     enable_xla_jit(False)
+
+
+def set_parallel_scaling(scaling_mode):
+    """Sets the scaling mode for data parallel acceleration.
+    Weak scaling involves increasing the problem size proportionally with the number of processors,
+    while strong scaling involves keeping the problem size fixed and increasing the number of processors.
+
+    Args:
+        scaling_mode (str): Whether 'weak' or 'strong'
+    """
+    global parallel_scaling
+    parallel_scaling = scaling_mode

deepxde/data/pde.py

@@ -100,6 +100,10 @@ def __init__(
                 raise ValueError(
                     "Parallel training via Horovod only supports pseudo train distribution."
                 )
+            if config.parallel_scaling == "strong":
+                raise ValueError(
+                    "Strong scaling is not supported with tensorflow.compat.v1. Please use weak scaling."
+                )
         self.anchors = None if anchors is None else anchors.astype(config.real(np))
         self.exclusions = exclusions