add mixed precision support to deepxde

deepxde/config.py

@@ -74,7 +74,7 @@ def set_default_float(value):
     The default floating point type is 'float32'.
 
     Args:
-        value (String): 'float16', 'float32', or 'float64'.
+        value (String): 'float16', 'float32', 'float64', or 'mixed' (mixed precision in https://arxiv.org/abs/2401.16645).
     """
     if value == "float16":
         print("Set the default float type to float16")
@@ -85,6 +85,20 @@ def set_default_float(value):
     elif value == "float64":
         print("Set the default float type to float64")
         real.set_float64()
+    elif value == "mixed":
+        print("Set the float type to mixed precision of float16 and float32")
+        real.set_mixed()
+        if backend_name == "tensorflow":
+            real.set_float16()
+            tf.keras.mixed_precision.set_global_policy("mixed_float16")
+            return # don't try to set it again below
+        if backend_name == "pytorch":
+            # Use float16 during the forward and backward passes, but store in float32
+            real.set_float32()
+        else:
+            raise ValueError(
+                f"{backend_name} backend does not currently support mixed precision."
+            )
     else:
         raise ValueError(f"{value} not supported in deepXDE")
     if backend_name in ["tensorflow.compat.v1", "tensorflow"]:

deepxde/model.py

@@ -357,9 +357,12 @@ def closure():
                 total_loss = torch.sum(losses)
                 self.opt.zero_grad()
                 total_loss.backward()
-                return total_loss
 
-            self.opt.step(closure)
+            def closure_mixed():
+                with torch.autocast(device_type="cuda", dtype=torch.float16):
+                    closure()
+
+            self.opt.step(closure if not config.real.mixed else closure_mixed)
             if self.lr_scheduler is not None:
                 self.lr_scheduler.step()
 

deepxde/real.py

@@ -7,6 +7,7 @@ class Real:
     def __init__(self, precision):
         self.precision = None
         self.reals = None
+        self.mixed = False
         if precision == 16:
             self.set_float16()
         elif precision == 32:
@@ -28,3 +29,6 @@ def set_float32(self):
     def set_float64(self):
         self.precision = 64
         self.reals = {np: np.float64, bkd.lib: bkd.float64}
+
+    def set_mixed(self):
+        self.mixed = True

docs/user/faq.rst

@@ -10,6 +10,8 @@ General usage
   | **A**: `#5`_
 - | **Q**: By default, DeepXDE uses ``float32``. How can I use ``float64``?
   | **A**: `#28`_
+- | **Q**: How can I use mixed precision training?
+  | **A**: Use ``dde.config.set_default_float("mixed")`` with the ``tensorflow`` or ``pytorch`` backends. See `this paper <https://doi.org/10.1016/j.cma.2024.117093>`_ for more information.
 - | **Q**: I want to set the global random seeds.
   | **A**: `#353`_
 - | **Q**: GPU.