feat: add ffn for ne and ea

Author: iProzd

deepmd/dpmodel/descriptor/dpa3.py

@@ -41,6 +41,9 @@ def __init__(
         d_sel: int = 10,
         d_rcut: float = 2.8,
         d_rcut_smth: float = 2.0,
+        use_ffn_node_edge_message: bool = False,
+        use_ffn_edge_angle_message: bool = False,
+        ffn_hidden_dim: int = 1024,
     ) -> None:
         r"""The constructor for the RepFlowArgs class which defines the parameters of the repflow block in DPA3 descriptor.
 
@@ -131,6 +134,9 @@ def __init__(
         self.d_sel = d_sel
         self.d_rcut = d_rcut
         self.d_rcut_smth = d_rcut_smth
+        self.use_ffn_node_edge_message = use_ffn_node_edge_message
+        self.use_ffn_edge_angle_message = use_ffn_edge_angle_message
+        self.ffn_hidden_dim = ffn_hidden_dim
 
     def __getitem__(self, key):
         if hasattr(self, key):

deepmd/pt/model/descriptor/dpa3.py

@@ -176,6 +176,9 @@ def init_subclass_params(sub_data, sub_class):
             d_sel=self.repflow_args.d_sel,
             d_rcut=self.repflow_args.d_rcut,
             d_rcut_smth=self.repflow_args.d_rcut_smth,
+            use_ffn_node_edge_message=self.repflow_args.use_ffn_node_edge_message,
+            use_ffn_edge_angle_message=self.repflow_args.use_ffn_edge_angle_message,
+            ffn_hidden_dim=self.repflow_args.ffn_hidden_dim,
             exclude_types=exclude_types,
             env_protection=env_protection,
             precision=precision,

deepmd/pt/model/descriptor/repflow_layer.py

@@ -17,6 +17,7 @@
     get_residual,
 )
 from deepmd.pt.model.network.mlp import (
+    FeedForward,
     MLPLayer,
 )
 from deepmd.pt.model.network.utils import (
@@ -63,6 +64,9 @@ def __init__(
         d_sel: int = 10,
         d_rcut: float = 2.8,
         d_rcut_smth: float = 2.0,
+        use_ffn_node_edge_message: bool = False,
+        use_ffn_edge_angle_message: bool = False,
+        ffn_hidden_dim: int = 1024,
         activation_function: str = "silu",
         update_style: str = "res_residual",
         update_residual: float = 0.1,
@@ -119,6 +123,11 @@ def __init__(
         self.d_rcut = d_rcut
         self.d_rcut_smth = d_rcut_smth
         self.dynamic_d_sel = (self.d_sel * 4) / self.sel_reduce_factor
+        self.use_ffn_node_edge_message = use_ffn_node_edge_message
+        self.use_ffn_edge_angle_message = use_ffn_edge_angle_message
+        self.ffn_hidden_dim = ffn_hidden_dim
+        if self.use_ffn_node_edge_message or self.use_ffn_edge_angle_message:
+            assert not self.optim_update, "FFN does not support optim update!"
 
         if self.update_dihedral:
             assert self.use_dynamic_sel, "Dihedral update requires dynamic selection!"
@@ -174,11 +183,20 @@ def __init__(
             )
 
         # node edge message
-        self.node_edge_linear = MLPLayer(
-            self.edge_info_dim,
-            self.n_multi_edge_message * n_dim,
-            precision=precision,
-            seed=child_seed(seed, 4),
+        self.node_edge_linear = (
+            MLPLayer(
+                self.edge_info_dim,
+                self.n_multi_edge_message * n_dim,
+                precision=precision,
+                seed=child_seed(seed, 4),
+            )
+            if not self.use_ffn_node_edge_message
+            else FeedForward(
+                self.edge_info_dim,
+                self.n_multi_edge_message * n_dim,
+                self.ffn_hidden_dim,
+                activation_function=self.activation_function,
+            )
         )
         if self.update_style == "res_residual":
             for head_index in range(self.n_multi_edge_message):
@@ -248,11 +266,20 @@ def __init__(
                     self.a_compress_e_linear = None
 
             # edge angle message
-            self.edge_angle_linear1 = MLPLayer(
-                self.angle_dim,
-                self.e_dim,
-                precision=precision,
-                seed=child_seed(seed, 10),
+            self.edge_angle_linear1 = (
+                MLPLayer(
+                    self.angle_dim,
+                    self.e_dim,
+                    precision=precision,
+                    seed=child_seed(seed, 10),
+                )
+                if not self.use_ffn_edge_angle_message
+                else FeedForward(
+                    self.angle_dim,
+                    self.e_dim,
+                    self.ffn_hidden_dim,
+                    activation_function=self.activation_function,
+                )
             )
             self.edge_angle_linear2 = MLPLayer(
                 self.e_dim,

deepmd/pt/model/descriptor/repflows.py

@@ -116,6 +116,9 @@ def __init__(
         d_sel: int = 10,
         d_rcut: float = 2.8,
         d_rcut_smth: float = 2.0,
+        use_ffn_node_edge_message: bool = False,
+        use_ffn_edge_angle_message: bool = False,
+        ffn_hidden_dim: int = 1024,
         optim_update: bool = True,
         seed: Optional[Union[int, list[int]]] = None,
     ) -> None:
@@ -248,6 +251,9 @@ def __init__(
         self.d_sel = d_sel
         self.d_rcut = d_rcut
         self.d_rcut_smth = d_rcut_smth
+        self.use_ffn_node_edge_message = use_ffn_node_edge_message
+        self.use_ffn_edge_angle_message = use_ffn_edge_angle_message
+        self.ffn_hidden_dim = ffn_hidden_dim
 
         self.n_dim = n_dim
         self.e_dim = e_dim
@@ -321,6 +327,9 @@ def __init__(
                     d_sel=self.d_sel,
                     d_rcut=self.d_rcut,
                     d_rcut_smth=self.d_rcut_smth,
+                    use_ffn_node_edge_message=self.use_ffn_node_edge_message,
+                    use_ffn_edge_angle_message=self.use_ffn_edge_angle_message,
+                    ffn_hidden_dim=self.ffn_hidden_dim,
                     seed=child_seed(child_seed(seed, 1), ii),
                 )
             )

deepmd/pt/model/network/mlp.py

@@ -40,6 +40,9 @@
     to_numpy_array,
     to_torch_tensor,
 )
+from deepmd.utils.version import (
+    check_version_compatibility,
+)
 
 
 def empty_t(shape, precision):
@@ -278,6 +281,83 @@ def check_load_param(ss):
         return obj
 
 
+class FeedForward(nn.Module):
+    """
+    A feed forward network with two linear layers and an activation function.
+    No dropout, no gate and no residual connection.
+    """
+
+    def __init__(
+        self,
+        num_in: int,
+        num_out: int,
+        hidden_dim: int,
+        activation_function: Optional[str] = None,
+        bias: bool = False,
+    ) -> None:
+        super().__init__()
+        self.num_in = num_in
+        self.num_out = num_out
+        self.hidden_dim = hidden_dim
+        self.activation_function = activation_function
+        self.bias = bias
+        self.w1 = MLPLayer(
+            num_in=num_in,
+            num_out=hidden_dim,
+            bias=bias,
+        )
+        self.act = ActivationFn(activation_function)
+        self.w2 = MLPLayer(
+            num_in=hidden_dim,
+            num_out=num_out,
+            bias=bias,
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.w2(self.act(self.w1(x)))
+
+    def serialize(self) -> dict:
+        """Serialize the networks to a dict.
+
+        Returns
+        -------
+        dict
+            The serialized networks.
+        """
+        data = {
+            "@class": "FeedForward",
+            "@version": 1,
+            "num_in": self.num_in,
+            "num_out": self.num_out,
+            "hidden_dim": self.hidden_dim,
+            "activation_function": self.activation_function,
+            "bias": self.bias,
+            "w1": self.w1.serialize(),
+            "w2": self.w2.serialize(),
+        }
+        return data
+
+    @classmethod
+    def deserialize(cls, data: dict) -> "FeedForward":
+        """Deserialize the networks from a dict.
+
+        Parameters
+        ----------
+        data : dict
+            The dict to deserialize from.
+        """
+        data = data.copy()
+        check_version_compatibility(data.pop("@version"), 1, 1)
+        data.pop("@class")
+        w1 = data.pop("w1")
+        w2 = data.pop("w2")
+
+        obj = cls(**data)
+        obj.w1 = MLPLayer.deserialize(w1)
+        obj.w2 = MLPLayer.deserialize(w2)
+        return obj
+
+
 MLP_ = make_multilayer_network(MLPLayer, nn.Module)
 
 

deepmd/utils/argcheck.py

@@ -1646,6 +1646,24 @@ def dpa3_repflow_args():
             optional=True,
             default=2.0,
         ),
+        Argument(
+            "use_ffn_node_edge_message",
+            bool,
+            optional=True,
+            default=False,
+        ),
+        Argument(
+            "use_ffn_edge_angle_message",
+            bool,
+            optional=True,
+            default=False,
+        ),
+        Argument(
+            "ffn_hidden_dim",
+            int,
+            optional=True,
+            default=1024,
+        ),
     ]