From afe33ee7aab6dfafc5d39e7ae4761ee0e4565aff Mon Sep 17 00:00:00 2001
From: Vladimir <vldud.off@gmail.com>
Date: Wed, 23 Oct 2024 18:38:08 +0300
Subject: [PATCH] Fix regularization loss

---
 deepxde/nn/tensorflow_compat_v1/deeponet.py | 54 +++++++++++++++------
 1 file changed, 40 insertions(+), 14 deletions(-)

diff --git a/deepxde/nn/tensorflow_compat_v1/deeponet.py b/deepxde/nn/tensorflow_compat_v1/deeponet.py
index d083dbae2..aa23b2263 100644
--- a/deepxde/nn/tensorflow_compat_v1/deeponet.py
+++ b/deepxde/nn/tensorflow_compat_v1/deeponet.py
@@ -422,14 +422,18 @@ def _dense(
         regularizer=None,
         trainable=True,
     ):
-        return tf.keras.layers.Dense(
+        dense = tf.keras.layers.Dense(
             units,
             activation=activation,
             use_bias=use_bias,
             kernel_initializer=self.kernel_initializer,
             kernel_regularizer=regularizer,
             trainable=trainable,
-        )(inputs)
+        )
+        out = dense(inputs)
+        if regularizer:
+            self.regularization_loss += tf.math.add_n(dense.losses)
+        return out
 
     def _stacked_dense(
         self, inputs, units, stack_size, activation=None, use_bias=True, trainable=True
@@ -636,23 +640,23 @@ def build_branch_net(self):
         else:
             # Fully connected network
             for i in range(1, len(self.layer_size_func) - 1):
-                y_func = tf.keras.layers.Dense(
+                y_func = self._dense(
+                    y_func,
                     self.layer_size_func[i],
                     activation=self.activation_branch,
-                    kernel_initializer=self.kernel_initializer,
-                    kernel_regularizer=self.regularizer,
-                )(y_func)
+                    regularizer=self.regularizer,
+                )
                 if self.dropout_rate_branch[i - 1] > 0:
                     y_func = tf.layers.dropout(
                         y_func,
                         rate=self.dropout_rate_branch[i - 1],
                         training=self.training,
                     )
-            y_func = tf.keras.layers.Dense(
+            y_func = self._dense(
+                y_func,
                 self.layer_size_func[-1],
-                kernel_initializer=self.kernel_initializer,
-                kernel_regularizer=self.regularizer,
-            )(y_func)
+                regularizer=self.regularizer,
+            )
         return y_func
 
     def build_trunk_net(self):
@@ -661,12 +665,12 @@ def build_trunk_net(self):
         if self._input_transform is not None:
             y_loc = self._input_transform(y_loc)
         for i in range(1, len(self.layer_size_loc)):
-            y_loc = tf.keras.layers.Dense(
+            y_loc = self._dense(
+                y_loc,
                 self.layer_size_loc[i],
                 activation=self.activation_trunk,
-                kernel_initializer=self.kernel_initializer,
-                kernel_regularizer=self.regularizer,
-            )(y_loc)
+                regularizer=self.regularizer,
+            )
             if self.dropout_rate_trunk[i - 1] > 0:
                 y_loc = tf.layers.dropout(
                     y_loc, rate=self.dropout_rate_trunk[i - 1], training=self.training
@@ -683,3 +687,25 @@ def merge_branch_trunk(self, branch, trunk):
     @staticmethod
     def concatenate_outputs(ys):
         return tf.stack(ys, axis=2)
+
+    def _dense(
+        self,
+        inputs,
+        units,
+        activation=None,
+        use_bias=True,
+        regularizer=None,
+        trainable=True,
+    ):
+        dense = tf.keras.layers.Dense(
+            units,
+            activation=activation,
+            use_bias=use_bias,
+            kernel_initializer=self.kernel_initializer,
+            kernel_regularizer=regularizer,
+            trainable=trainable,
+        )
+        out = dense(inputs)
+        if regularizer:
+            self.regularization_loss += tf.math.add_n(dense.losses)
+        return out