fix

Author: HydrogenSulfate

deepmd/pd/entrypoints/main.py

@@ -353,6 +353,12 @@ def freeze(
     )
     model = inference.Tester(model, head=head).model
     model.eval()
+    # print(model.get_buffer_rcut.__func__.__qualname__)
+    # print(model.get_buffer_rcut.__func__.__module__)
+    # print(model.get_buffer_rcut.__func__.__code__.co_filename)
+    # print(model.get_buffer_rcut.__func__.__code__.co_firstlineno)
+    # print(model.get_buffer_type_map())
+    # exit()
     from paddle.static import (
         InputSpec,
     )
@@ -411,14 +417,14 @@ def freeze(
             full_graph=True,
         )
     for method_name in [
-        "get_rcut",
-        "get_type_map",
-        "get_dim_fparam",
-        "get_dim_aparam",
-        "get_intensive",
-        "get_sel_type",
-        "get_numb_dos",
-        "get_task_dim",
+        "get_buffer_rcut",
+        "get_buffer_type_map",
+        "get_buffer_dim_fparam",
+        "get_buffer_dim_aparam",
+        "get_buffer_intensive",
+        "get_buffer_sel_type",
+        "get_buffer_numb_dos",
+        "get_buffer_task_dim",
     ]:
         if hasattr(model, method_name):
             setattr(

deepmd/pd/infer/deep_eval.py

@@ -150,9 +150,9 @@ def __init__(
             self.static_model = False
         elif str(self.model_path).endswith(".json"):
             self.dp = paddle.jit.load(self.model_path.split(".json")[0])
-            self.rcut = self.dp.get_rcut().item()
+            self.rcut = self.dp.get_buffer_rcut().item()
             self.type_map: list[str] = "".join(
-                [chr(x) for x in self.dp.get_type_map().numpy()]
+                [chr(x) for x in self.dp.get_buffer_type_map().numpy()]
             ).split(" ")
             config = paddle_inference.Config(
                 self.model_path, self.model_path.replace(".json", ".pdiparams")
@@ -214,13 +214,13 @@ def get_type_map(self) -> list[str]:
     def get_dim_fparam(self) -> int:
         """Get the number (dimension) of frame parameters of this DP."""
         if self.static_model:
-            return self.dp.get_dim_fparam()
+            return self.dp.get_buffer_dim_fparam()
         return self.dp.model["Default"].get_dim_fparam()
 
     def get_dim_aparam(self) -> int:
         """Get the number (dimension) of atomic parameters of this DP."""
         if self.static_model:
-            return self.dp.get_dim_aparam()
+            return self.dp.get_buffer_dim_aparam()
         return self.dp.model["Default"].get_dim_aparam()
 
     def has_default_fparam(self) -> bool:

deepmd/pd/model/atomic_model/base_atomic_model.py

@@ -134,6 +134,10 @@ def get_type_map(self) -> list[str]:
         """Get the type map."""
         return self.type_map
 
+    def get_buffer_type_map(self) -> list[str]:
+        """Get the type map."""
+        return self.buffer_type_map
+
     def get_compute_stats_distinguish_types(self) -> bool:
         """Get whether the fitting net computes stats which are not distinguished between different types of atoms."""
         return True
@@ -581,8 +585,9 @@ def _store_out_stat(
         paddle.assign(out_std_data, self.out_std)
 
     def get_ntypes(self):
-        if paddle.in_dynamic_mode():
-            return len(self.type_map)
+        return len(self.type_map)
+
+    def get_buffer_ntypes(self):
         return self.buffer_ntypes
 
     def _fetch_out_stat(

deepmd/pd/model/atomic_model/dp_atomic_model.py

@@ -71,11 +71,14 @@ def __init__(
         def _string_to_array(s: Union[str, list[str]]) -> list[int]:
             return [ord(c) for c in s]
 
-        self.register_buffer(
-            "buffer_type_map",
-            paddle.to_tensor(_string_to_array(" ".join(self.type_map)), dtype="int32"),
-        )
-        self.buffer_type_map.name = "buffer_type_map"
+        if type_map is not None:
+            self.register_buffer(
+                "buffer_type_map",
+                paddle.to_tensor(
+                    _string_to_array(" ".join(self.type_map)), dtype="int32"
+                ),
+            )
+            self.buffer_type_map.name = "buffer_type_map"
         if hasattr(self.descriptor, "has_message_passing"):
             # register 'has_message_passing' as buffer(cast to int32 as problems may meets with vector<bool>)
             self.register_buffer(
@@ -148,15 +151,23 @@ def fitting_output_def(self) -> FittingOutputDef:
 
     def get_rcut(self) -> float:
         """Get the cut-off radius."""
-        if paddle.in_dynamic_mode():
-            return self.rcut
-        return self.descriptor.get_rcut()
+        return self.rcut
 
     def get_sel(self) -> list[int]:
         """Get the neighbor selection."""
-        if paddle.in_dynamic_mode():
-            return self.sel
-        return self.descriptor.get_sel()
+        return self.sel
+
+    def get_buffer_type_map(self) -> paddle.Tensor:
+        """Get the neighbor selection."""
+        return self.buffer_type_map
+
+    def get_buffer_rcut(self) -> paddle.Tensor:
+        """Get the cut-off radius."""
+        return self.descriptor.get_buffer_rcut()
+
+    def get_buffer_sel(self) -> paddle.Tensor:
+        """Get the neighbor selection."""
+        return self.descriptor.get_buffer_sel()
 
     def set_case_embd(self, case_idx: int):
         """

deepmd/pd/model/descriptor/dpa1.py

@@ -325,10 +325,18 @@ def get_rcut(self) -> float:
         """Returns the cut-off radius."""
         return self.se_atten.get_rcut()
 
+    def get_buffer_rcut(self) -> float:
+        """Returns the cut-off radius."""
+        return self.se_atten.get_buffer_rcut()
+
     def get_rcut_smth(self) -> float:
         """Returns the radius where the neighbor information starts to smoothly decay to 0."""
         return self.se_atten.get_rcut_smth()
 
+    def get_buffer_rcut_smth(self) -> float:
+        """Returns the radius where the neighbor information starts to smoothly decay to 0."""
+        return self.se_atten.get_buffer_rcut_smth()
+
     def get_nsel(self) -> int:
         """Returns the number of selected atoms in the cut-off radius."""
         return self.se_atten.get_nsel()
@@ -343,8 +351,10 @@ def get_ntypes(self) -> int:
 
     def get_type_map(self) -> list[str]:
         """Get the name to each type of atoms."""
-        if paddle.in_dynamic_mode():
-            return self.type_map
+        return self.type_map
+
+    def get_buffer_type_map(self) -> paddle.Tensor:
+        """Get the name to each type of atoms."""
         return self.buffer_type_map
 
     def get_dim_out(self) -> int:

deepmd/pd/model/descriptor/dpa2.py

@@ -334,15 +334,19 @@ def init_subclass_params(sub_data, sub_class):
 
     def get_rcut(self) -> float:
         """Returns the cut-off radius."""
-        if paddle.in_dynamic_mode():
-            return self.rcut
-        return self.repinit.get_rcut()
+        return self.rcut
 
     def get_rcut_smth(self) -> float:
         """Returns the radius where the neighbor information starts to smoothly decay to 0."""
-        if paddle.in_dynamic_mode():
-            return self.rcut_smth
-        return self.repinit.get_rcut_smth()
+        return self.rcut_smth
+
+    def get_buffer_rcut(self) -> paddle.Tensor:
+        """Returns the cut-off radius."""
+        return self.repinit.get_buffer_rcut()
+
+    def get_buffer_rcut_smth(self) -> paddle.Tensor:
+        """Returns the radius where the neighbor information starts to smoothly decay to 0."""
+        return self.repinit.get_buffer_rcut_smth()
 
     def get_nsel(self) -> int:
         """Returns the number of selected atoms in the cut-off radius."""
@@ -780,7 +784,7 @@ def forward(
             type_embedding = None
         g1, _, _, _, _ = self.repinit(
             nlist_dict[
-                get_multiple_nlist_key(self.repinit.get_rcut(), self.repinit.get_nsel())
+                get_multiple_nlist_key(self.repinit.rcut, sum(self.repinit.sel))
             ],
             extended_coord,
             extended_atype,

deepmd/pd/model/descriptor/dpa3.py

@@ -176,10 +176,11 @@ def init_subclass_params(sub_data, sub_class):
         self.use_loc_mapping = use_loc_mapping
         self.use_tebd_bias = use_tebd_bias
         self.type_map = type_map
-        self.register_buffer(
-            "buffer_type_map",
-            paddle.to_tensor([ord(c) for c in " ".join(self.type_map)]),
-        )
+        if type_map is not None:
+            self.register_buffer(
+                "buffer_type_map",
+                paddle.to_tensor([ord(c) for c in " ".join(self.type_map)]),
+            )
         self.tebd_dim = self.repflow_args.n_dim
         self.type_embedding = TypeEmbedNet(
             ntypes,
@@ -222,24 +223,30 @@ def init_subclass_params(sub_data, sub_class):
 
     def get_rcut(self) -> float:
         """Returns the cut-off radius."""
-        if paddle.in_dynamic_mode():
-            return self.rcut
-        return self.repflows.get_rcut()
+        return self.rcut
 
     def get_rcut_smth(self) -> float:
         """Returns the radius where the neighbor information starts to smoothly decay to 0."""
-        if paddle.in_dynamic_mode():
-            return self.rcut_smth
-        return self.repflows.get_rcut_smth()
+        return self.rcut_smth
+
+    def get_buffer_rcut(self) -> paddle.Tensor:
+        """Returns the cut-off radius."""
+        return self.repflows.get_buffer_rcut()
+
+    def get_buffer_rcut_smth(self) -> paddle.Tensor:
+        """Returns the radius where the neighbor information starts to smoothly decay to 0."""
+        return self.repflows.get_buffer_rcut_smth()
 
     def get_nsel(self) -> int:
         """Returns the number of selected atoms in the cut-off radius."""
         return sum(self.sel)
 
     def get_sel(self) -> list[int]:
         """Returns the number of selected atoms for each type."""
-        if paddle.in_dynamic_mode():
-            return self.sel
+        return self.sel
+
+    def get_buffer_sel(self) -> paddle.Tensor:
+        """Returns the number of selected atoms for each type."""
         return self.repflows.get_sel()
 
     def get_ntypes(self) -> int:
@@ -248,8 +255,10 @@ def get_ntypes(self) -> int:
 
     def get_type_map(self) -> list[str]:
         """Get the name to each type of atoms."""
-        if paddle.in_dynamic_mode():
-            return self.type_map
+        return self.type_map
+
+    def get_buffer_type_map(self) -> paddle.Tensor:
+        """Get the name to each type of atoms."""
         return self.buffer_type_map
 
     def get_dim_out(self) -> int:

deepmd/pd/model/descriptor/repflows.py

@@ -358,8 +358,10 @@ def __init__(
 
     def get_rcut(self) -> float:
         """Returns the cut-off radius."""
-        if paddle.in_dynamic_mode():
-            return self.e_rcut
+        return self.e_rcut
+
+    def get_buffer_rcut(self) -> paddle.Tensor:
+        """Returns the cut-off radius."""
         return self.buffer_rcut
 
     def get_rcut_smth(self) -> float:
@@ -372,8 +374,10 @@ def get_nsel(self) -> int:
 
     def get_sel(self) -> list[int]:
         """Returns the number of selected atoms for each type."""
-        if paddle.in_dynamic_mode():
-            return self.sel
+        return self.sel
+
+    def get_buffer_sel(self) -> paddle.Tensor:
+        """Returns the number of selected atoms for each type."""
         return self.buffer_sel
 
     def get_ntypes(self) -> int:

deepmd/pd/model/descriptor/repformers.py

@@ -323,14 +323,18 @@ def __init__(
 
     def get_rcut(self) -> float:
         """Returns the cut-off radius."""
-        if paddle.in_dynamic_mode():
-            return self.rcut
-        return self.buffer_rcut
+        return self.rcut
 
     def get_rcut_smth(self) -> float:
         """Returns the radius where the neighbor information starts to smoothly decay to 0."""
-        if paddle.in_dynamic_mode():
-            return self.rcut_smth
+        return self.rcut_smth
+
+    def get_buffer_rcut(self) -> paddle.Tensor:
+        """Returns the cut-off radius."""
+        return self.buffer_rcut
+
+    def get_buffer_rcut_smth(self) -> paddle.Tensor:
+        """Returns the radius where the neighbor information starts to smoothly decay to 0."""
         return self.buffer_rcut_smth
 
     def get_nsel(self) -> int:

deepmd/pd/model/descriptor/se_a.py

@@ -141,8 +141,10 @@ def get_ntypes(self) -> int:
 
     def get_type_map(self) -> list[str]:
         """Get the name to each type of atoms."""
-        if paddle.in_dynamic_mode():
-            return self.type_map
+        return self.type_map
+
+    def get_buffer_type_map(self) -> paddle.Tensor:
+        """Get the name to each type of atoms."""
         return self.buffer_type_map
 
     def get_dim_out(self) -> int:
@@ -519,14 +521,18 @@ def __init__(
 
     def get_rcut(self) -> float:
         """Returns the cut-off radius."""
-        if paddle.in_dynamic_mode():
-            return self.rcut
-        return self.buffer_rcut
+        return self.rcut
 
     def get_rcut_smth(self) -> float:
         """Returns the radius where the neighbor information starts to smoothly decay to 0."""
-        if paddle.in_dynamic_mode():
-            return self.rcut_smth
+        return self.rcut_smth
+
+    def get_buffer_rcut(self) -> paddle.Tensor:
+        """Returns the cut-off radius."""
+        return self.buffer_rcut
+
+    def get_buffer_rcut_smth(self) -> paddle.Tensor:
+        """Returns the radius where the neighbor information starts to smoothly decay to 0."""
         return self.buffer_rcut_smth
 
     def get_nsel(self) -> int: