Revert "pd: support dpa1 (deepmodeling#4414)"

This reverts commit e8167ce.

Author: HydrogenSulfate

deepmd/pd/entrypoints/main.py

@@ -230,7 +230,7 @@ def train(
     use_pretrain_script: bool = False,
     force_load: bool = False,
     output: str = "out.json",
-) -> None:
+):
     log.info("Configuration path: %s", input_file)
     SummaryPrinter()()
     with open(input_file) as fin:
@@ -321,26 +321,18 @@ def train(
     # save min_nbor_dist
     if min_nbor_dist is not None:
         if not multi_task:
-            trainer.model.min_nbor_dist = paddle.to_tensor(
-                min_nbor_dist,
-                dtype=paddle.float64,
-                place=DEVICE,
-            )
+            trainer.model.min_nbor_dist = min_nbor_dist
         else:
             for model_item in min_nbor_dist:
-                trainer.model[model_item].min_nbor_dist = paddle.to_tensor(
-                    min_nbor_dist[model_item],
-                    dtype=paddle.float64,
-                    place=DEVICE,
-                )
+                trainer.model[model_item].min_nbor_dist = min_nbor_dist[model_item]
     trainer.run()
 
 
 def freeze(
     model: str,
     output: str = "frozen_model.json",
     head: Optional[str] = None,
-) -> None:
+):
     paddle.set_flags(
         {
             "FLAGS_save_cf_stack_op": 1,
@@ -391,7 +383,7 @@ def change_bias(
     numb_batch: int = 0,
     model_branch: Optional[str] = None,
     output: Optional[str] = None,
-) -> None:
+):
     if input_file.endswith(".pd"):
         old_state_dict = paddle.load(input_file)
         model_state_dict = copy.deepcopy(old_state_dict.get("model", old_state_dict))

deepmd/pd/loss/ener.py

@@ -10,6 +10,7 @@
     TaskLoss,
 )
 from deepmd.pd.utils import (
+    decomp,
     env,
 )
 from deepmd.pd.utils.env import (
@@ -223,7 +224,10 @@ def forward(self, input_dict, model, label, natoms, learning_rate, mae=False):
 
             if self.relative_f is not None:
                 force_label_3 = force_label.reshape([-1, 3])
-                norm_f = force_label_3.norm(axis=1, keepdim=True) + self.relative_f
+                # norm_f = force_label_3.norm(axis=1, keepdim=True) + self.relative_f
+                norm_f = (
+                    decomp.norm(force_label_3, axis=1, keepdim=True) + self.relative_f
+                )
                 diff_f_3 = diff_f.reshape([-1, 3])
                 diff_f_3 = diff_f_3 / norm_f
                 diff_f = diff_f_3.reshape([-1])

deepmd/pd/model/atomic_model/dp_atomic_model.py

@@ -1,4 +1,5 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
+import copy
 import functools
 import logging
 from typing import (
@@ -51,7 +52,7 @@ def __init__(
         fitting,
         type_map: list[str],
         **kwargs,
-    ) -> None:
+    ):
         super().__init__(type_map, **kwargs)
         ntypes = len(type_map)
         self.type_map = type_map
@@ -200,7 +201,7 @@ def serialize(self) -> dict:
 
     @classmethod
     def deserialize(cls, data) -> "DPAtomicModel":
-        data = data.copy()
+        data = copy.deepcopy(data)
         check_version_compatibility(data.pop("@version", 1), 2, 1)
         data.pop("@class", None)
         data.pop("type", None)
@@ -211,37 +212,6 @@ def deserialize(cls, data) -> "DPAtomicModel":
         obj = super().deserialize(data)
         return obj
 
-    def enable_compression(
-        self,
-        min_nbor_dist: float,
-        table_extrapolate: float = 5,
-        table_stride_1: float = 0.01,
-        table_stride_2: float = 0.1,
-        check_frequency: int = -1,
-    ) -> None:
-        """Call descriptor enable_compression().
-
-        Parameters
-        ----------
-        min_nbor_dist
-            The nearest distance between atoms
-        table_extrapolate
-            The scale of model extrapolation
-        table_stride_1
-            The uniform stride of the first table
-        table_stride_2
-            The uniform stride of the second table
-        check_frequency
-            The overflow check frequency
-        """
-        self.descriptor.enable_compression(
-            min_nbor_dist,
-            table_extrapolate,
-            table_stride_1,
-            table_stride_2,
-            check_frequency,
-        )
-
     def forward_atomic(
         self,
         extended_coord,
@@ -308,7 +278,7 @@ def compute_or_load_stat(
         self,
         sampled_func,
         stat_file_path: Optional[DPPath] = None,
-    ) -> None:
+    ):
         """
         Compute or load the statistics parameters of the model,
         such as mean and standard deviation of descriptors or the energy bias of the fitting net.

deepmd/pd/model/descriptor/__init__.py

@@ -5,10 +5,6 @@
 from .descriptor import (
     DescriptorBlock,
 )
-from .dpa1 import (
-    DescrptBlockSeAtten,
-    DescrptDPA1,
-)
 from .env_mat import (
     prod_env_mat,
 )
@@ -21,8 +17,6 @@
     "BaseDescriptor",
     "DescriptorBlock",
     "DescrptBlockSeA",
-    "DescrptBlockSeAtten",
-    "DescrptDPA1",
     "DescrptSeA",
     "prod_env_mat",
 ]