Towards SoftAdapt loss balancing for tf.compat.v1

deepxde/callbacks.py

@@ -571,3 +571,34 @@ def on_epoch_end(self):
             raise ValueError(
                 "`num_bcs` changed! Please update the loss function by `model.compile`."
             )
+
+
+class SoftAdapt(Callback):
+    """Use adaptive loss balancing.
+
+    Args:
+        beta: If beta > 0, then softAdapt will pay more attention the worst performing
+            loss component. If beta < 0, then SoftAdapt will assign higher weights
+            to the better performing components. Beta==0 is the trivial case and
+            all loss components will have coefficient 1.
+        epsilon: parameter to prevent overflows.
+
+    """
+
+    def __init__(self, beta=.1, epsilon=1e-8):
+        super().__init__()
+
+        self.beta = beta
+        self.epsilon = epsilon
+
+    def on_train_begin(self):
+        loss_weights = tf.constant(self.model.loss_weights)
+        loss_weights = dde.Variable(loss_weights, trainable=False, dtype=loss_weights.dtype)
+        loss_weights *= 0
+
+        self.model.loss_weights = loss_weights
+
+        print(loss_weights, 'loss_weights')
+        # Allow instances to be re-used.
+        # Evaluate coefficients.
+        # Update weights.