remove eval_ hooks

Author: Han Wang

deepmd/dpmodel/atomic_model/dp_atomic_model.py

@@ -3,8 +3,6 @@
     Any,
 )
 
-import array_api_compat
-
 from deepmd.dpmodel.array_api import (
     Array,
 )
@@ -56,10 +54,6 @@ def __init__(
         if hasattr(self.fitting, "reinit_exclude"):
             self.fitting.reinit_exclude(self.atom_exclude_types)
         self.type_map = type_map
-        self.enable_eval_descriptor_hook = False
-        self.enable_eval_fitting_last_layer_hook = False
-        self.eval_descriptor_list: list[Array] = []
-        self.eval_fitting_last_layer_list: list[Array] = []
         super().init_out_stat()
 
     def fitting_output_def(self) -> FittingOutputDef:
@@ -132,37 +126,6 @@ def enable_compression(
             check_frequency,
         )
 
-    def set_eval_descriptor_hook(self, enable: bool) -> None:
-        """Set the hook for evaluating descriptor and clear the cache."""
-        self.enable_eval_descriptor_hook = enable
-        self.eval_descriptor_list.clear()
-
-    def eval_descriptor(self) -> Array:
-        """Evaluate the descriptor by concatenating cached results."""
-        if not self.eval_descriptor_list:
-            raise RuntimeError(
-                "eval_descriptor_list is empty. "
-                "Call set_eval_descriptor_hook(True) and perform a forward pass first."
-            )
-        xp = array_api_compat.array_namespace(self.eval_descriptor_list[0])
-        return xp.concat(self.eval_descriptor_list, axis=0)
-
-    def set_eval_fitting_last_layer_hook(self, enable: bool) -> None:
-        """Set the hook for evaluating fitting last layer output and clear the cache."""
-        self.enable_eval_fitting_last_layer_hook = enable
-        self.fitting.set_return_middle_output(enable)
-        self.eval_fitting_last_layer_list.clear()
-
-    def eval_fitting_last_layer(self) -> Array:
-        """Evaluate the fitting last layer output by concatenating cached results."""
-        if not self.eval_fitting_last_layer_list:
-            raise RuntimeError(
-                "eval_fitting_last_layer_list is empty. "
-                "Call set_eval_fitting_last_layer_hook(True) and perform a forward pass first."
-            )
-        xp = array_api_compat.array_namespace(self.eval_fitting_last_layer_list[0])
-        return xp.concat(self.eval_fitting_last_layer_list, axis=0)
-
     def forward_atomic(
         self,
         extended_coord: Array,
@@ -203,8 +166,6 @@ def forward_atomic(
             nlist,
             mapping=mapping,
         )
-        if self.enable_eval_descriptor_hook:
-            self.eval_descriptor_list.append(descriptor)
         ret = self.fitting(
             descriptor,
             atype,
@@ -214,11 +175,6 @@ def forward_atomic(
             fparam=fparam,
             aparam=aparam,
         )
-        if self.enable_eval_fitting_last_layer_hook:
-            assert "middle_output" in ret, (
-                "eval_fitting_last_layer not supported for this fitting net!"
-            )
-            self.eval_fitting_last_layer_list.append(ret.pop("middle_output"))
         return ret
 
     def change_type_map(

deepmd/dpmodel/model/dp_model.py

@@ -1,9 +1,6 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
 
 
-from deepmd.dpmodel.array_api import (
-    Array,
-)
 from deepmd.dpmodel.descriptor.base_descriptor import (
     BaseDescriptor,
 )
@@ -55,19 +52,3 @@ def get_fitting_net(self) -> BaseFitting:
     def get_descriptor(self) -> BaseDescriptor:
         """Get the descriptor."""
         return self.atomic_model.descriptor
-
-    def set_eval_descriptor_hook(self, enable: bool) -> None:
-        """Set the hook for evaluating descriptor."""
-        self.atomic_model.set_eval_descriptor_hook(enable)
-
-    def eval_descriptor(self) -> Array:
-        """Evaluate the descriptor."""
-        return self.atomic_model.eval_descriptor()
-
-    def set_eval_fitting_last_layer_hook(self, enable: bool) -> None:
-        """Set the hook for evaluating fitting last layer output."""
-        self.atomic_model.set_eval_fitting_last_layer_hook(enable)
-
-    def eval_fitting_last_layer(self) -> Array:
-        """Evaluate the fitting last layer output."""
-        return self.atomic_model.eval_fitting_last_layer()

source/tests/consistent/model/test_ener.py

@@ -757,54 +757,6 @@ def test_forward_common_lower_alias(self) -> None:
         for key in ret_call:
             np.testing.assert_equal(ret_call[key], ret_fc[key])
 
-    def test_eval_descriptor(self) -> None:
-        """eval_descriptor should produce consistent results across dp and pt."""
-        # dpmodel
-        self.dp_model.set_eval_descriptor_hook(True)
-        self.dp_model.call_lower(
-            self.extended_coord,
-            self.extended_atype,
-            self.nlist,
-            self.mapping,
-        )
-        dp_desc = self.dp_model.eval_descriptor()
-
-        # pt
-        self.pt_model.set_eval_descriptor_hook(True)
-        self.pt_model.forward_common_lower(
-            numpy_to_torch(self.extended_coord),
-            numpy_to_torch(self.extended_atype),
-            numpy_to_torch(self.nlist),
-            numpy_to_torch(self.mapping),
-        )
-        pt_desc = torch_to_numpy(self.pt_model.eval_descriptor())
-
-        np.testing.assert_allclose(dp_desc, pt_desc, rtol=1e-10, atol=1e-10)
-
-    def test_eval_fitting_last_layer(self) -> None:
-        """eval_fitting_last_layer should produce consistent results across dp and pt."""
-        # dpmodel
-        self.dp_model.set_eval_fitting_last_layer_hook(True)
-        self.dp_model.call_lower(
-            self.extended_coord,
-            self.extended_atype,
-            self.nlist,
-            self.mapping,
-        )
-        dp_fl = self.dp_model.eval_fitting_last_layer()
-
-        # pt
-        self.pt_model.set_eval_fitting_last_layer_hook(True)
-        self.pt_model.forward_common_lower(
-            numpy_to_torch(self.extended_coord),
-            numpy_to_torch(self.extended_atype),
-            numpy_to_torch(self.nlist),
-            numpy_to_torch(self.mapping),
-        )
-        pt_fl = torch_to_numpy(self.pt_model.eval_fitting_last_layer())
-
-        np.testing.assert_allclose(dp_fl, pt_fl, rtol=1e-10, atol=1e-10)
-
     def test_model_output_def(self) -> None:
         """model_output_def should return the same keys and shapes on dp and pt."""
         dp_def = self.dp_model.model_output_def().get_data()