add update_use_layernorm

Author: iProzd

deepmd/dpmodel/descriptor/dpa3.py

@@ -180,6 +180,7 @@ def __init__(
         use_exp_switch: bool = False,
         use_dynamic_sel: bool = False,
         sel_reduce_factor: float = 10.0,
+        update_use_layernorm: bool = False,
     ) -> None:
         self.n_dim = n_dim
         self.e_dim = e_dim
@@ -210,6 +211,7 @@ def __init__(
         self.use_exp_switch = use_exp_switch
         self.use_dynamic_sel = use_dynamic_sel
         self.sel_reduce_factor = sel_reduce_factor
+        self.update_use_layernorm = update_use_layernorm
 
     def __getitem__(self, key: str) -> Any:
         if hasattr(self, key):

deepmd/pt/model/descriptor/dpa3.py

@@ -166,6 +166,7 @@ def init_subclass_params(sub_data: Any, sub_class: Any) -> Any:
             use_dynamic_sel=self.repflow_args.use_dynamic_sel,
             sel_reduce_factor=self.repflow_args.sel_reduce_factor,
             use_loc_mapping=use_loc_mapping,
+            update_use_layernorm=self.repflow_args.update_use_layernorm,
             exclude_types=exclude_types,
             env_protection=env_protection,
             precision=precision,

deepmd/pt/model/descriptor/repflow_layer.py

@@ -58,6 +58,7 @@ def __init__(
         use_dynamic_sel: bool = False,
         sel_reduce_factor: float = 10.0,
         smooth_edge_update: bool = False,
+        update_use_layernorm: bool = False,
         activation_function: str = "silu",
         update_style: str = "res_residual",
         update_residual: float = 0.1,
@@ -97,6 +98,7 @@ def __init__(
         self.update_style = update_style
         self.update_residual = update_residual
         self.update_residual_init = update_residual_init
+        self.update_use_layernorm = update_use_layernorm
         self.a_compress_e_rate = a_compress_e_rate
         self.a_compress_use_split = a_compress_use_split
         self.precision = precision
@@ -203,6 +205,17 @@ def __init__(
                 )
             )
 
+        if self.update_use_layernorm:
+            self.node_layernorm = torch.nn.LayerNorm(self.n_dim)
+            self.edge_layernorm = torch.nn.LayerNorm(self.e_dim)
+            self.angle_layernorm = (
+                torch.nn.LayerNorm(self.a_dim) if self.update_angle else None
+            )
+        else:
+            self.node_layernorm = None
+            self.edge_layernorm = None
+            self.angle_layernorm = None
+
         if self.update_angle:
             self.angle_dim = self.a_dim
             if self.a_compress_rate == 0:
@@ -1133,6 +1146,14 @@ def forward(
 
         # update angle_ebd
         a_updated = self.list_update(a_update_list, "angle")
+        if self.update_use_layernorm:
+            assert self.node_layernorm is not None
+            n_updated = self.node_layernorm(n_updated)
+            assert self.edge_layernorm is not None
+            e_updated = self.edge_layernorm(e_updated)
+            if self.update_angle:
+                assert self.angle_layernorm is not None
+                a_updated = self.angle_layernorm(a_updated)
         return n_updated, e_updated, a_updated
 
     @torch.jit.export

deepmd/pt/model/descriptor/repflows.py

@@ -220,6 +220,7 @@ def __init__(
         use_dynamic_sel: bool = False,
         sel_reduce_factor: float = 10.0,
         use_loc_mapping: bool = True,
+        update_use_layernorm: bool = False,
         optim_update: bool = True,
         seed: Optional[Union[int, list[int]]] = None,
         trainable: bool = True,
@@ -285,6 +286,7 @@ def __init__(
         self.precision = precision
         self.epsilon = 1e-4
         self.seed = seed
+        self.update_use_layernorm = update_use_layernorm
 
         self.edge_embd = MLPLayer(
             1,
@@ -330,6 +332,7 @@ def __init__(
                     use_dynamic_sel=self.use_dynamic_sel,
                     sel_reduce_factor=self.sel_reduce_factor,
                     smooth_edge_update=self.smooth_edge_update,
+                    update_use_layernorm=self.update_use_layernorm,
                     seed=child_seed(child_seed(seed, 1), ii),
                     trainable=trainable,
                 )

deepmd/utils/argcheck.py

@@ -1667,6 +1667,12 @@ def dpa3_repflow_args() -> list[Argument]:
             default=10.0,
             doc=doc_sel_reduce_factor,
         ),
+        Argument(
+            "update_use_layernorm",
+            bool,
+            optional=True,
+            default=False,
+        ),
     ]