add mixed precision support to deepxde

Author: g-w1

deepxde/config.py

@@ -79,6 +79,19 @@ def set_default_float(value):
     if value == "float16":
         print("Set the default float type to float16")
         real.set_float16()
+    elif value == "mixed":
+        print("Set training policy to mixed")
+        real.set_mixed()
+        if backend_name == "tensorflow":
+            real.set_float16()
+            tf.keras.mixed_precision.set_global_policy("mixed_float16")
+        elif backend_name == "pytorch":
+            # we cast to float16 during the passes in the training loop, but store in float32
+            real.set_float32()
+        else:
+            raise ValueError(
+                f"{backend_name} backend does not currently support mixed precision in deepXDE"
+            )
     elif value == "float32":
         print("Set the default float type to float32")
         real.set_float32()

deepxde/model.py

@@ -353,10 +353,20 @@ def outputs_losses_test(inputs, targets, auxiliary_vars):
 
         def train_step(inputs, targets, auxiliary_vars):
             def closure():
-                losses = outputs_losses_train(inputs, targets, auxiliary_vars)[1]
-                total_loss = torch.sum(losses)
-                self.opt.zero_grad()
-                total_loss.backward()
+                if config.real.mixed:
+                    with torch.autocast(device_type="cuda", dtype=torch.float16):
+                        losses = outputs_losses_train(inputs, targets, auxiliary_vars)[
+                            1
+                        ]
+                        total_loss = torch.sum(losses)
+                        # we do the backprop in float16
+                        self.opt.zero_grad()
+                        total_loss.backward()
+                else:
+                    losses = outputs_losses_train(inputs, targets, auxiliary_vars)[1]
+                    total_loss = torch.sum(losses)
+                    self.opt.zero_grad()
+                    total_loss.backward()
                 return total_loss
 
             self.opt.step(closure)

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:
@@ -17,6 +18,9 @@ def __init__(self, precision):
     def __call__(self, package):
         return self.reals[package]
 
+    def set_mixed(self):
+        self.mixed = True
+
     def set_float16(self):
         self.precision = 16
         self.reals = {np: np.float16, bkd.lib: bkd.float16}