Manual loss weights adaptation in TF2.0

deepxde/callbacks.py

@@ -571,3 +571,50 @@ def on_epoch_end(self):
             raise ValueError(
                 "`num_bcs` changed! Please update the loss function by `model.compile`."
             )
+
+
+class PrintLossWeight(Callback):
+    def __init__(self, period):
+        super().__init__()
+        """Print the loss weights every period epochs.
+
+        Args:
+            period: Interval (number of epochs) between printing loss weights.
+
+        """
+        self.period = period
+        self.initial_loss_weights = None
+        self.current_loss_weights = None
+
+    def on_epoch_begin(self):
+        if self.model.train_state.epoch == 0:
+            self.initial_loss_weights = self.model.loss_weights.numpy().tolist()
+        else:
+            self.current_loss_weights = self.model.loss_weights.numpy().tolist()
+        if self.model.train_state.epoch % self.period == 0:
+            print("Initial loss weights:", self.initial_loss_weights)
+            print("Current loss weights:", self.current_loss_weights)
+
+
+class ManualDynamicLossWeight(Callback):
+    def __init__(self, epoch2change, value, loss_idx):
+        super().__init__()
+        """Change the loss weights at a specific epoch.
+
+        Args:
+            epoch2change: The epoch at which to change the loss weight
+            value: The value to change the loss weight to
+            idx: The index of the loss weight to change
+
+        """
+        self.epoch2change = epoch2change
+        self.value = value
+        self.loss_idx = loss_idx
+
+    def on_epoch_begin(self):
+        if self.model.train_state.epoch == self.epoch2change:
+            current_loss_weights = self.model.loss_weights.numpy()
+            current_loss_weights[self.loss_idx] = self.value
+            self.model.loss_weights = tf.convert_to_tensor(
+                current_loss_weights, dtype=config.default_float()
+            )

deepxde/model.py

@@ -119,7 +119,9 @@ def compile(
             print("Compiling model...")
         self.opt_name = optimizer
         loss_fn = losses_module.get(loss)
-        self.loss_weights = loss_weights
+        self.loss_weights = tf.convert_to_tensor(
+            loss_weights, dtype=config.default_float()
+        )
         if external_trainable_variables is None:
             self.external_trainable_variables = []
         else:
@@ -202,7 +204,9 @@ def _compile_tensorflow(self, lr, loss_fn, decay):
         def outputs(training, inputs):
             return self.net(inputs, training=training)
 
-        def outputs_losses(training, inputs, targets, auxiliary_vars, losses_fn):
+        def outputs_losses(
+            training, inputs, targets, auxiliary_vars, losses_fn, loss_weights
+        ):
             self.net.auxiliary_vars = auxiliary_vars
             # Don't call outputs() decorated by @tf.function above, otherwise the
             # gradient of outputs wrt inputs will be lost here.
@@ -218,29 +222,41 @@ def outputs_losses(training, inputs, targets, auxiliary_vars, losses_fn):
                 losses += [tf.math.reduce_sum(self.net.losses)]
             losses = tf.convert_to_tensor(losses)
             # Weighted losses
-            if self.loss_weights is not None:
-                losses *= self.loss_weights
+            if loss_weights is not None:
+                losses *= loss_weights
             return outputs_, losses
 
         @tf.function(jit_compile=config.xla_jit)
-        def outputs_losses_train(inputs, targets, auxiliary_vars):
+        def outputs_losses_train(inputs, targets, auxiliary_vars, loss_weights):
             return outputs_losses(
-                True, inputs, targets, auxiliary_vars, self.data.losses_train
+                True,
+                inputs,
+                targets,
+                auxiliary_vars,
+                self.data.losses_train,
+                loss_weights,
             )
 
         @tf.function(jit_compile=config.xla_jit)
-        def outputs_losses_test(inputs, targets, auxiliary_vars):
+        def outputs_losses_test(inputs, targets, auxiliary_vars, loss_weights):
             return outputs_losses(
-                False, inputs, targets, auxiliary_vars, self.data.losses_test
+                False,
+                inputs,
+                targets,
+                auxiliary_vars,
+                self.data.losses_test,
+                loss_weights,
             )
 
         opt = optimizers.get(self.opt_name, learning_rate=lr, decay=decay)
 
         @tf.function(jit_compile=config.xla_jit)
-        def train_step(inputs, targets, auxiliary_vars):
+        def train_step(inputs, targets, auxiliary_vars, loss_weights):
             # inputs and targets are np.ndarray and automatically converted to Tensor.
             with tf.GradientTape() as tape:
-                losses = outputs_losses_train(inputs, targets, auxiliary_vars)[1]
+                losses = outputs_losses_train(
+                    inputs, targets, auxiliary_vars, loss_weights
+                )[1]
                 total_loss = tf.math.reduce_sum(losses)
             trainable_variables = (
                 self.net.trainable_variables + self.external_trainable_variables
@@ -531,7 +547,7 @@ def _outputs(self, training, inputs):
             outs = self.outputs(self.net.params, training, inputs)
         return utils.to_numpy(outs)
 
-    def _outputs_losses(self, training, inputs, targets, auxiliary_vars):
+    def _outputs_losses(self, training, inputs, targets, auxiliary_vars, loss_weights):
         if training:
             outputs_losses = self.outputs_losses_train
         else:
@@ -540,7 +556,7 @@ def _outputs_losses(self, training, inputs, targets, auxiliary_vars):
             feed_dict = self.net.feed_dict(training, inputs, targets, auxiliary_vars)
             return self.sess.run(outputs_losses, feed_dict=feed_dict)
         if backend_name == "tensorflow":
-            outs = outputs_losses(inputs, targets, auxiliary_vars)
+            outs = outputs_losses(inputs, targets, auxiliary_vars, loss_weights)
         elif backend_name == "pytorch":
             self.net.requires_grad_(requires_grad=False)
             outs = outputs_losses(inputs, targets, auxiliary_vars)
@@ -552,12 +568,12 @@ def _outputs_losses(self, training, inputs, targets, auxiliary_vars):
             outs = outputs_losses(inputs, targets, auxiliary_vars)
         return utils.to_numpy(outs[0]), utils.to_numpy(outs[1])
 
-    def _train_step(self, inputs, targets, auxiliary_vars):
+    def _train_step(self, inputs, targets, auxiliary_vars, loss_weights):
         if backend_name == "tensorflow.compat.v1":
             feed_dict = self.net.feed_dict(True, inputs, targets, auxiliary_vars)
             self.sess.run(self.train_step, feed_dict=feed_dict)
         elif backend_name in ["tensorflow", "paddle"]:
-            self.train_step(inputs, targets, auxiliary_vars)
+            self.train_step(inputs, targets, auxiliary_vars, loss_weights)
         elif backend_name == "pytorch":
             self.train_step(inputs, targets, auxiliary_vars)
         elif backend_name == "jax":
@@ -669,6 +685,7 @@ def _train_sgd(self, iterations, display_every):
                 self.train_state.X_train,
                 self.train_state.y_train,
                 self.train_state.train_aux_vars,
+                self.loss_weights,
             )
 
             self.train_state.epoch += 1
@@ -827,12 +844,14 @@ def _test(self):
             self.train_state.X_train,
             self.train_state.y_train,
             self.train_state.train_aux_vars,
+            self.loss_weights,
         )
         self.train_state.y_pred_test, self.train_state.loss_test = self._outputs_losses(
             False,
             self.train_state.X_test,
             self.train_state.y_test,
             self.train_state.test_aux_vars,
+            self.loss_weights,
         )
 
         if isinstance(self.train_state.y_test, (list, tuple)):

examples/pinn_inverse/elliptic_inverse_field_manual_dynamic_loss_weights.py

@@ -0,0 +1,81 @@
+"""Backend supported: tensorflow.compat.v1, tensorflow, pytorch, paddle"""
+
+# import sys
+import deepxde as dde
+import matplotlib.pyplot as plt
+import numpy as np
+from deepxde.callbacks import PrintLossWeight, ManualDynamicLossWeight
+
+dde.config.disable_xla_jit()
+from deepxde.backend import set_default_backend
+
+set_default_backend("tensorflow")
+
+
+def gen_traindata(num):
+    # generate num equally-spaced points from -1 to 1
+    xvals = np.linspace(-1, 1, num).reshape(num, 1)
+    uvals = np.sin(np.pi * xvals)
+    return xvals, uvals
+
+
+def pde(x, y):
+    u, q = y[:, 0:1], y[:, 1:2]
+    du_xx = dde.grad.hessian(y, x, component=0, i=0, j=0)
+    return -du_xx + q
+
+
+def sol(x):
+    # solution is u(x) = sin(pi*x), q(x) = -pi^2 * sin(pi*x)
+    return np.sin(np.pi * x)
+
+
+geom = dde.geometry.Interval(-1, 1)
+bc = dde.icbc.DirichletBC(geom, sol, lambda _, on_boundary: on_boundary, component=0)
+ob_x, ob_u = gen_traindata(100)
+observe_u = dde.icbc.PointSetBC(ob_x, ob_u, component=0)
+
+data = dde.data.PDE(
+    geom,
+    pde,
+    [bc, observe_u],
+    num_domain=200,
+    num_boundary=2,
+    anchors=ob_x,
+    num_test=1000,
+)
+
+net = dde.nn.FNN([1, 40, 40, 40, 2], "tanh", "Glorot uniform")
+PrintLossWeight_cb = PrintLossWeight(period=1)
+ManualDynamicLossWeight_cb = ManualDynamicLossWeight(
+    epoch2change=5000, value=1, loss_idx=0
+)
+model = dde.Model(data, net)
+model.compile("adam", lr=0.0001, loss_weights=[0, 100, 1000])
+losshistory, train_state = model.train(
+    iterations=20000,
+    display_every=1,
+    callbacks=[PrintLossWeight_cb, ManualDynamicLossWeight_cb],
+)
+# dde.saveplot(losshistory, train_state, issave=True, isplot=True)
+
+# view results
+x = geom.uniform_points(500)
+yhat = model.predict(x)
+uhat, qhat = yhat[:, 0:1], yhat[:, 1:2]
+
+utrue = np.sin(np.pi * x)
+print("l2 relative error for u: " + str(dde.metrics.l2_relative_error(utrue, uhat)))
+plt.figure()
+plt.plot(x, utrue, "-", label="u_true")
+plt.plot(x, uhat, "--", label="u_NN")
+plt.legend()
+
+qtrue = -np.pi**2 * np.sin(np.pi * x)
+print("l2 relative error for q: " + str(dde.metrics.l2_relative_error(qtrue, qhat)))
+plt.figure()
+plt.plot(x, qtrue, "-", label="q_true")
+plt.plot(x, qhat, "--", label="q_NN")
+plt.legend()
+
+plt.show()