add mixed precision support to deepxde

Author: g-w1

deepxde/config.py

@@ -71,10 +71,10 @@ def default_float():
 def set_default_float(value):
     """Sets the default float type.
 
-    The default floating point type is 'float32'.
+    The default floating point type is 'float32'. Mixed precision uses the method in the paper: `J. Hayford, J. Goldman-Wetzler, E. Wang, & L. Lu. Speeding up and reducing memory usage for scientific machine learning via mixed precision. Computer Methods in Applied Mechanics and Engineering, 428, 117093, 2024 <https://doi.org/10.1016/j.cma.2024.117093>`_.
 
     Args:
-        value (String): 'float16', 'float32', or 'float64'.
+        value (String): 'float16', 'float32', 'float64', or 'mixed' (mixed precision).
     """
     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

@@ -372,9 +372,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