Backend paddle: add mfnn net; add regularizer; add optimizers;

Author: lijialin03

deepxde/backend/backend.py

@@ -502,3 +502,12 @@ def sparse_dense_matmul(x, y):
     Returns:
         Tensor: The multiplication result.
     """
+
+def l1_decay(x):
+    """Implement the L1 weight decay regularization."""
+
+def l2_decay(x):
+    """Implement the L2 weight decay regularization."""
+
+def l1_l2_decay(x,y):
+    """Implement the L1 and L2 weight decay regularization."""

deepxde/backend/paddle/tensor.py

@@ -229,3 +229,9 @@ def matmul(x, y):
 
 def sparse_dense_matmul(x, y):
     return paddle.sparse.matmul(x, y)
+
+def l1_decay(x):
+    return paddle.regularizer.L1Decay(coeff=x)
+
+def l2_decay(x):
+    return paddle.regularizer.L2Decay(coeff=x)

deepxde/backend/tensorflow_compat_v1/tensor.py

@@ -245,3 +245,12 @@ def matmul(x, y):
 
 def sparse_dense_matmul(x, y):
     return tf.sparse.sparse_dense_matmul(x, y)
+
+def l1_decay(x):
+    return tf.keras.regularizers.L1(l1=x)
+
+def l2_decay(x):
+    return tf.keras.regularizers.L2(l2=x)
+
+def l1_l2_decay(x,y):
+    return tf.keras.regularizers.L1L2(l1=x, l2=y)

deepxde/data/mf.py

@@ -1,7 +1,8 @@
 import numpy as np
 
 from .data import Data
-from ..backend import tf
+from .. import backend as bkd
+from .. import config
 from ..utils import run_if_any_none, standardize
 
 
@@ -116,7 +117,7 @@ def losses_train(self, targets, outputs, loss_fn, inputs, model, aux=None):
         return [loss_lo, loss_hi]
 
     def losses_test(self, targets, outputs, loss_fn, inputs, model, aux=None):
-        return [0, loss_fn(targets[1], outputs[1])]
+        return [bkd.as_tensor(0, dtype=config.real(bkd.lib)), loss_fn(targets[1], outputs[1])]
 
     @run_if_any_none("X_train", "y_train")
     def train_next_batch(self, batch_size=None):

deepxde/model.py

@@ -506,8 +506,13 @@ def outputs_losses_test(inputs, targets, auxiliary_vars):
         trainable_variables = (
             list(self.net.parameters()) + self.external_trainable_variables
         )
+        regularizer = getattr(self.net, 'regularizer', None)
+        if regularizer is not None:
+            weight_decay = self.net.regularizer_value if self.opt_name == "adamw" else self.net.regularizer
+        else: 
+            weight_decay = None
         self.opt = optimizers.get(
-            trainable_variables, self.opt_name, learning_rate=lr, decay=decay
+            trainable_variables, self.opt_name, learning_rate=lr, decay=decay, weight_decay=weight_decay,
         )
 
         def train_step(inputs, targets, auxiliary_vars):

deepxde/nn/paddle/__init__.py

@@ -4,6 +4,7 @@
     "DeepONet",
     "DeepONetCartesianProd",
     "FNN",
+    "MfNN",
     "MsFFN",
     "PFNN",
     "STMsFFN",
@@ -12,3 +13,4 @@
 from .deeponet import DeepONet, DeepONetCartesianProd
 from .fnn import FNN, PFNN
 from .msffn import MsFFN, STMsFFN
+from .mfnn import MfNN

deepxde/nn/paddle/fnn.py

@@ -3,12 +3,20 @@
 from .nn import NN
 from .. import activations
 from .. import initializers
+from .. import regularizers
 
 
 class FNN(NN):
     """Fully-connected neural network."""
 
-    def __init__(self, layer_sizes, activation, kernel_initializer):
+    def __init__(
+        self,
+        layer_sizes,
+        activation,
+        kernel_initializer,
+        regularization=None,
+        dropout_rate=0.0,
+    ):
         super().__init__()
         if isinstance(activation, list):
             if not (len(layer_sizes) - 1) == len(activation):
@@ -20,12 +28,16 @@ def __init__(self, layer_sizes, activation, kernel_initializer):
             self.activation = activations.get(activation)
         initializer = initializers.get(kernel_initializer)
         initializer_zero = initializers.get("zeros")
+        self.regularizer = regularizers.get(regularization)
+        self.regularizer_value = regularization[1:] if regularization is not None else None
+        self.dropout_rate = dropout_rate
 
         self.linears = paddle.nn.LayerList()
         for i in range(1, len(layer_sizes)):
             self.linears.append(paddle.nn.Linear(layer_sizes[i - 1], layer_sizes[i]))
             initializer(self.linears[-1].weight)
             initializer_zero(self.linears[-1].bias)
+        self.dropout = paddle.nn.Dropout(p=dropout_rate) if dropout_rate > 0.0 else None
 
     def forward(self, inputs):
         x = inputs
@@ -37,6 +49,8 @@ def forward(self, inputs):
                 if isinstance(self.activation, list)
                 else self.activation(linear(x))
             )
+        if self.dropout is not None:
+            x = self.dropout(x)
         x = self.linears[-1](x)
         if self._output_transform is not None:
             x = self._output_transform(inputs, x)

deepxde/nn/paddle/mfnn.py

@@ -0,0 +1,119 @@
+import paddle
+
+from .nn import NN
+from .. import activations
+from .. import initializers
+from .. import regularizers
+from ... import config
+
+
+class MfNN(NN):
+    """Multifidelity neural networks."""
+
+    def __init__(
+        self,
+        layer_sizes_low_fidelity,
+        layer_sizes_high_fidelity,
+        activation,
+        kernel_initializer,
+        regularization=None,
+        residue=False,
+        trainable_low_fidelity=True,
+        trainable_high_fidelity=True,
+    ):
+        super().__init__()
+        self.layer_size_lo = layer_sizes_low_fidelity
+        self.layer_size_hi = layer_sizes_high_fidelity
+
+        self.activation = activations.get(activation)
+        self.activation_tanh = activations.get("tanh")
+        self.initializer = initializers.get(kernel_initializer)
+        self.initializer_zero = initializers.get("zeros")
+        self.trainable_lo = trainable_low_fidelity
+        self.trainable_hi = trainable_high_fidelity
+        self.residue = residue
+        self.regularizer = regularizers.get(regularization)
+        self.regularizer_value = regularization[1:] if regularization is not None else None
+        
+        # low fidelity
+        self.linears_lo = self.init_dense(self.layer_size_lo, self.trainable_lo)
+        
+        # high fidelity
+        # linear part
+        self.linears_hi_l = paddle.nn.Linear(
+            in_features=self.layer_size_lo[0] + self.layer_size_lo[-1],
+            out_features=self.layer_size_hi[-1],
+            weight_attr=paddle.ParamAttr(initializer=self.initializer),
+            bias_attr=paddle.ParamAttr(initializer=self.initializer_zero),
+        )
+        if not self.trainable_hi:
+            for param in self.linears_hi_l.parameters():
+                param.stop_gradient = False
+        # nonlinear part
+        self.layer_size_hi = [self.layer_size_lo[0] + self.layer_size_lo[-1]] + self.layer_size_hi
+        self.linears_hi = self.init_dense(self.layer_size_hi, self.trainable_hi)
+        # linear + nonlinear
+        if not self.residue:
+            alpha = self.init_alpha(0.0, self.trainable_hi)
+            self.add_parameter("alpha",alpha)
+        else:
+            alpha1 = self.init_alpha(0.0, self.trainable_hi)
+            alpha2 = self.init_alpha(0.0, self.trainable_hi)
+            self.add_parameter("alpha1",alpha1)
+            self.add_parameter("alpha2",alpha2)
+    
+    def init_dense(self, layer_size, trainable):
+        linears = paddle.nn.LayerList()
+        for i in range(len(layer_size) - 1):
+            linear = paddle.nn.Linear(
+                in_features=layer_size[i],
+                out_features=layer_size[i + 1],
+                weight_attr=paddle.ParamAttr(initializer=self.initializer),
+                bias_attr=paddle.ParamAttr(initializer=self.initializer_zero),
+            )
+            if not trainable:
+                for param in linear.parameters():
+                    param.stop_gradient = False
+            linears.append(linear)
+        return linears
+
+    def init_alpha(self, value, trainable):
+        alpha = paddle.create_parameter(
+            shape=[1], 
+            dtype=config.real(paddle), 
+            default_initializer=paddle.nn.initializer.Constant(value),
+        )
+        alpha.stop_gradient=not trainable
+        return alpha
+
+    def forward(self, inputs):
+        x = inputs.astype(config.real(paddle))
+        # low fidelity
+        y = x
+        for i, linear in enumerate(self.linears_lo):
+            y = linear(y)
+            if i != len(self.linears_lo) - 1:
+                y = self.activation(y)
+        y_lo = y
+
+        # high fidelity
+        x_hi = paddle.concat([x, y_lo], axis=1)
+        # linear
+        y_hi_l = self.linears_hi_l(x_hi)
+        # nonlinear
+        y = x_hi
+        for i, linear in enumerate(self.linears_hi):
+            y = linear(y)
+            if i != len(self.linears_hi) - 1:
+                y = self.activation(y)
+        y_hi_nl = y
+        # linear + nonlinear
+        if not self.residue:
+            alpha = self.activation_tanh(self.alpha)
+            y_hi = y_hi_l + alpha * y_hi_nl
+        else:
+            alpha1 = self.activation_tanh(self.alpha1)
+            alpha2 = self.activation_tanh(self.alpha2)
+            y_hi = y_lo + 0.1 * (alpha1 * y_hi_l + alpha2 * y_hi_nl)
+
+        return y_lo, y_hi

deepxde/nn/regularizers.py

@@ -1,4 +1,5 @@
-from ..backend import tf
+from .. import backend as bkd
+from ..backend import backend_name
 
 
 def get(identifier):
@@ -22,12 +23,15 @@ def get(identifier):
     if not factor:
         raise ValueError("Regularization factor must be provided.")
 
-    if name == "l1":
-        return tf.keras.regularizers.L1(l1=factor[0])
-    if name == "l2":
-        return tf.keras.regularizers.L2(l2=factor[0])
-    if name in ("l1l2", "l1+l2"):
-        if len(factor) < 2:
-            raise ValueError("L1L2 regularizer requires both L1/L2 penalties.")
-        return tf.keras.regularizers.L1L2(l1=factor[0], l2=factor[1])
+    try:
+        if name == "l1":
+            return bkd.l1_decay(factor[0])
+        if name == "l2":
+            return bkd.l2_decay(factor[0])
+        if name in ("l1l2", "l1+l2"):
+            # TODO: only supported by 'tensorflow.compat.v1' now.
+            if len(factor) < 2:
+                return bkd.l1_l2_decay(factor[0], factor[1])
+    except Exception:
+        print(f"{name} regularization to be implemented for backend {backend_name} now.")
     raise ValueError(f"Unknown regularizer name: {name}")

deepxde/optimizers/paddle/optimizers.py

@@ -19,12 +19,14 @@ def is_external_optimizer(optimizer):
     return optimizer in ["L-BFGS", "L-BFGS-B"]
 
 
-def get(params, optimizer, learning_rate=None, decay=None):
+def get(params, optimizer, learning_rate=None, decay=None, weight_decay=None):
     """Retrieves an Optimizer instance."""
     if isinstance(optimizer, paddle.optimizer.Optimizer):
         return optimizer
 
     if optimizer in ["L-BFGS", "L-BFGS-B"]:
+        if weight_decay is not None:
+            raise ValueError("L-BFGS optimizer doesn't support weight_decay")
         if learning_rate is not None or decay is not None:
             print("Warning: learning rate is ignored for {}".format(optimizer))
         optim = paddle.optimizer.LBFGS(
@@ -46,5 +48,17 @@ def get(params, optimizer, learning_rate=None, decay=None):
         learning_rate = _get_lr_scheduler(learning_rate, decay)
 
     if optimizer == "adam":
-        return paddle.optimizer.Adam(learning_rate=learning_rate, parameters=params)
+        return paddle.optimizer.Adam(learning_rate=learning_rate, parameters=params, weight_decay=weight_decay)
+    elif optimizer == "sgd":
+        return paddle.optimizer.SGD(learning_rate=learning_rate, parameters=params, weight_decay=weight_decay)
+    elif optimizer == "rmsprop":
+        return paddle.optimizer.RMSProp(
+            learning_rate=learning_rate, parameters=params, weight_decay=weight_decay,
+        )
+    elif optimizer == "adamw":
+        if weight_decay[0] == 0:
+            raise ValueError("AdamW optimizer requires non-zero weight decay")
+        return paddle.optimizer.AdamW(
+            learning_rate=learning_rate, parameters=params, weight_decay=weight_decay[0],
+        )
     raise NotImplementedError(f"{optimizer} to be implemented for backend Paddle.")