[Refine] Refine granularity of arch module

docs/zh/api/arch.md

@@ -5,14 +5,23 @@
     options:
       members:
         - Arch
+        - FullyConnectedLayer
+        - DeepOperatorLayer
+        - LorenzEmbeddingLayer
+        - RosslerEmbeddingLayer
+        - CylinderEmbeddingLayer
+        - DiscriminatorLayer
+        - PhysformerGPT2Layer
+        - GeneratorLayer
+        - UNetExLayer
         - MLP
         - DeepONet
-        - DeepPhyLSTM
         - LorenzEmbedding
         - RosslerEmbedding
         - CylinderEmbedding
         - Generator
         - Discriminator
+        - DeepPhyLSTM
         - PhysformerGPT2
         - ModelList
         - AFNONet

ppsci/arch/__init__.py

@@ -16,36 +16,54 @@
 
 from ppsci.arch.base import Arch  # isort:skip
 from ppsci.arch.mlp import MLP  # isort:skip
+from ppsci.arch.mlp import FullyConnectedLayer  # isort:skip
 from ppsci.arch.deeponet import DeepONet  # isort:skip
+from ppsci.arch.deeponet import DeepOperatorLayer  # isort:skip
 from ppsci.arch.embedding_koopman import LorenzEmbedding  # isort:skip
+from ppsci.arch.embedding_koopman import LorenzEmbeddingLayer  # isort:skip
 from ppsci.arch.embedding_koopman import RosslerEmbedding  # isort:skip
+from ppsci.arch.embedding_koopman import RosslerEmbeddingLayer  # isort:skip
 from ppsci.arch.embedding_koopman import CylinderEmbedding  # isort:skip
+from ppsci.arch.embedding_koopman import CylinderEmbeddingLayer  # isort:skip
 from ppsci.arch.gan import Generator  # isort:skip
+from ppsci.arch.gan import GeneratorLayer  # isort:skip
 from ppsci.arch.gan import Discriminator  # isort:skip
+from ppsci.arch.gan import DiscriminatorLayer  # isort:skip
 from ppsci.arch.phylstm import DeepPhyLSTM  # isort:skip
 from ppsci.arch.physx_transformer import PhysformerGPT2  # isort:skip
+from ppsci.arch.physx_transformer import PhysformerGPT2Layer  # isort:skip
 from ppsci.arch.model_list import ModelList  # isort:skip
 from ppsci.arch.afno import AFNONet  # isort:skip
 from ppsci.arch.afno import PrecipNet  # isort:skip
 from ppsci.arch.unetex import UNetEx  # isort:skip
+from ppsci.arch.unetex import UNetExLayer  # isort:skip
 from ppsci.utils import logger  # isort:skip
 
 
 __all__ = [
     "Arch",
     "MLP",
+    "FullyConnectedLayer",
     "DeepONet",
-    "DeepPhyLSTM",
+    "DeepOperatorLayer",
     "LorenzEmbedding",
+    "LorenzEmbeddingLayer",
     "RosslerEmbedding",
+    "RosslerEmbeddingLayer",
     "CylinderEmbedding",
+    "CylinderEmbeddingLayer",
     "Generator",
+    "GeneratorLayer",
     "Discriminator",
+    "DiscriminatorLayer",
+    "DeepPhyLSTM",
     "PhysformerGPT2",
+    "PhysformerGPT2Layer",
     "ModelList",
     "AFNONet",
     "PrecipNet",
     "UNetEx",
+    "UNetExLayer",
     "build_model",
 ]
 

ppsci/arch/activation.py

@@ -24,6 +24,10 @@
 from ppsci.utils import initializer
 from ppsci.utils import misc
 
+__all__ = [
+    "get_activation",
+]
+
 
 class Stan(nn.Layer):
     """Self-scalable Tanh.

ppsci/arch/afno.py

@@ -30,6 +30,11 @@
 from ppsci.arch import base
 from ppsci.utils import initializer
 
+__all__ = [
+    "AFNONet",
+    "PrecipNet",
+]
+
 
 def drop_path(
     x: paddle.Tensor,

ppsci/arch/base.py

@@ -24,6 +24,10 @@
 
 from ppsci.utils import logger
 
+__all__ = [
+    "Arch",
+]
+
 
 class Arch(nn.Layer):
     """Base class for Network."""

ppsci/arch/deeponet.py

@@ -24,16 +24,19 @@
 from ppsci.arch import base
 from ppsci.arch import mlp
 
+__all__ = [
+    "DeepOperatorLayer",
+    "DeepONet",
+]
 
-class DeepONet(base.Arch):
-    """Deep operator network.
+
+class DeepOperatorLayer(base.Arch):
+    """Deep operator network, core implementation of `DeepONet`.
 
     [Lu et al. Learning nonlinear operators via DeepONet based on the universal approximation theorem of operators. Nat Mach Intell, 2021.](https://doi.org/10.1038/s42256-021-00302-5)
 
     Args:
-        u_key (str): Name of function data for input function u(x).
-        y_key (str): Name of location data for input function G(u).
-        G_key (str): Output name of predicted G(u)(y).
+        trunck_dim (int): Dimension of sampled u(x)(1 for scalar function, >1 for vector function).
         num_loc (int): Number of sampled u(x), i.e. `m` in paper.
         num_features (int): Number of features extracted from u(x), same for y.
         branch_num_layers (int): Number of hidden layers of branch net.
@@ -52,8 +55,8 @@ class DeepONet(base.Arch):
 
     Examples:
         >>> import ppsci
-        >>> model = ppsci.arch.DeepONet(
-        ...     "u", "y", "G",
+        >>> model = ppsci.arch.DeepOperatorLayer(
+        ...     1,
         ...     100, 40,
         ...     1, 1,
         ...     40, 40,
@@ -64,9 +67,7 @@ class DeepONet(base.Arch):
 
     def __init__(
         self,
-        u_key: str,
-        y_key: str,
-        G_key: str,
+        trunck_dim: int,
         num_loc: int,
         num_features: int,
         branch_num_layers: int,
@@ -82,33 +83,25 @@ def __init__(
         use_bias: bool = True,
     ):
         super().__init__()
-        self.u_key = u_key
-        self.y_key = y_key
-        self.input_keys = (u_key, y_key)
-        self.output_keys = (G_key,)
-
-        self.branch_net = mlp.MLP(
-            (self.u_key,),
-            ("b",),
+        self.trunck_dim = trunck_dim
+        self.branch_net = mlp.FullyConnectedLayer(
+            num_loc,
+            num_features,
             branch_num_layers,
             branch_hidden_size,
             branch_activation,
             branch_skip_connection,
             branch_weight_norm,
-            input_dim=num_loc,
-            output_dim=num_features,
         )
 
-        self.trunk_net = mlp.MLP(
-            (self.y_key,),
-            ("t",),
+        self.trunk_net = mlp.FullyConnectedLayer(
+            trunck_dim,
+            num_features,
             trunk_num_layers,
             trunk_hidden_size,
             trunk_activation,
             trunk_skip_connection,
             trunk_weight_norm,
-            input_dim=1,
-            output_dim=num_features,
         )
         self.trunk_act = act_mod.get_activation(trunk_activation)
 
@@ -120,28 +113,112 @@ def __init__(
                 attr=nn.initializer.Constant(0.0),
             )
 
-    def forward(self, x):
-        if self._input_transform is not None:
-            x = self._input_transform(x)
-
+    def forward(self, u, y):
         # Branch net to encode the input function
-        u_features = self.branch_net(x)[self.branch_net.output_keys[0]]
+        u_features = self.branch_net(u)
 
         # Trunk net to encode the domain of the output function
-        y_features = self.trunk_net(x)
-        y_features = self.trunk_act(y_features[self.trunk_net.output_keys[0]])
+        y_features = self.trunk_net(y)
+        y_features = self.trunk_act(y_features)
 
         # Dot product
         G_u = paddle.einsum("bi,bi->b", u_features, y_features)  # [batch_size, ]
-        G_u = paddle.reshape(G_u, [-1, 1])  # reshape [batch_size, ] to [batch_size, 1]
+        G_u = paddle.reshape(
+            G_u, [-1, self.trunck_dim]
+        )  # reshape [batch_size, ] to [batch_size, 1]
 
         # Add bias
         if self.use_bias:
             G_u += self.b
 
-        result_dict = {
-            self.output_keys[0]: G_u,
-        }
+        return G_u
+
+
+class DeepONet(DeepOperatorLayer):
+    """Deep operator network.
+    Different from `DeepOperatorLayer`, this class accepts input/output string key(s) for symbolic computation.
+
+    [Lu et al. Learning nonlinear operators via DeepONet based on the universal approximation theorem of operators. Nat Mach Intell, 2021.](https://doi.org/10.1038/s42256-021-00302-5)
+
+    Args:
+        u_key (str): Name of function data for input function u(x).
+        y_key (str): Name of location data for input function G(u).
+        G_key (str): Output name of predicted G(u)(y).
+        num_loc (int): Number of sampled u(x), i.e. `m` in paper.
+        num_features (int): Number of features extracted from u(x), same for y.
+        branch_num_layers (int): Number of hidden layers of branch net.
+        trunk_num_layers (int): Number of hidden layers of trunk net.
+        branch_hidden_size (Union[int, Tuple[int, ...]]): Number of hidden size of branch net.
+            An integer for all layers, or list of integer specify each layer's size.
+        trunk_hidden_size (Union[int, Tuple[int, ...]]): Number of hidden size of trunk net.
+            An integer for all layers, or list of integer specify each layer's size.
+        branch_skip_connection (bool, optional): Whether to use skip connection for branch net. Defaults to False.
+        trunk_skip_connection (bool, optional): Whether to use skip connection for trunk net. Defaults to False.
+        branch_activation (str, optional): Name of activation function. Defaults to "tanh".
+        trunk_activation (str, optional): Name of activation function. Defaults to "tanh".
+        branch_weight_norm (bool, optional): Whether to apply weight norm on parameter(s) for branch net. Defaults to False.
+        trunk_weight_norm (bool, optional): Whether to apply weight norm on parameter(s) for trunk net. Defaults to False.
+        use_bias (bool, optional): Whether to add bias on predicted G(u)(y). Defaults to True.
+
+    Examples:
+        >>> import ppsci
+        >>> model = ppsci.arch.DeepONet(
+        ...     "u", "y", "G",
+        ...     100, 40,
+        ...     1, 1,
+        ...     40, 40,
+        ...     branch_activation="relu", trunk_activation="relu",
+        ...     use_bias=True,
+        ... )
+    """
+
+    def __init__(
+        self,
+        u_key: str,
+        y_key: str,
+        G_key: str,
+        num_loc: int,
+        num_features: int,
+        branch_num_layers: int,
+        trunk_num_layers: int,
+        branch_hidden_size: Union[int, Tuple[int, ...]],
+        trunk_hidden_size: Union[int, Tuple[int, ...]],
+        branch_skip_connection: bool = False,
+        trunk_skip_connection: bool = False,
+        branch_activation: str = "tanh",
+        trunk_activation: str = "tanh",
+        branch_weight_norm: bool = False,
+        trunk_weight_norm: bool = False,
+        use_bias: bool = True,
+    ):
+        super().__init__()
+        self.input_keys = (u_key, y_key)
+        self.output_keys = (G_key,)
+
+        super().__init__(
+            1,
+            num_loc,
+            num_features,
+            branch_num_layers,
+            trunk_num_layers,
+            branch_hidden_size,
+            trunk_hidden_size,
+            branch_skip_connection,
+            trunk_skip_connection,
+            branch_activation,
+            trunk_activation,
+            branch_weight_norm,
+            trunk_weight_norm,
+            use_bias,
+        )
+
+    def forward(self, x):
+        if self._input_transform is not None:
+            x = self._input_transform(x)
+
+        G_u = super().forward(x[self.input_keys[0]], x[self.input_keys[1]])
+        result_dict = {self.output_keys[0]: G_u}
+
         if self._output_transform is not None:
             result_dict = self._output_transform(x, result_dict)