add DPModifier

Author: ChiahsinChu

deepmd/pt/modifier/dp_modifier.py

@@ -0,0 +1,92 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+
+import torch
+
+from deepmd.pt.model.model import (
+    BaseModel,
+)
+from deepmd.pt.modifier.base_modifier import (
+    BaseModifier,
+)
+from deepmd.pt.utils import (
+    env,
+)
+from deepmd.pt.utils.serialization import (
+    serialize_from_file,
+)
+
+
+class DPModifier(BaseModifier):
+    def __init__(
+        self,
+        dp_model: torch.nn.Module | None = None,
+        dp_model_file_name: str | None = None,
+        use_cache: bool = True,
+        **kwargs,
+    ) -> None:
+        """Constructor."""
+        super().__init__(use_cache=use_cache)
+
+        if dp_model_file_name is None and dp_model is None:
+            raise AttributeError("`model_name` or `model` should be specified.")
+        if dp_model_file_name is not None and dp_model is not None:
+            raise AttributeError(
+                "`model_name` and `model` cannot be used simultaneously."
+            )
+
+        if dp_model is not None:
+            self._model = dp_model.to(env.DEVICE)
+        if dp_model_file_name is not None:
+            data = serialize_from_file(dp_model_file_name)
+            assert data["model"]["type"] == "standard"
+            self._model = (
+                BaseModel.get_class_by_type(data["model"]["fitting"]["type"])
+                .deserialize(data["model"])
+                .to(env.DEVICE)
+            )
+        self._model.eval()
+        # use jit model for inference
+        self.model = torch.jit.script(self._model)
+
+    def serialize(self) -> dict:
+        """Serialize the modifier.
+
+        Returns
+        -------
+        dict
+            The serialized data
+        """
+        dd = BaseModifier.serialize(self)
+        dd.update(
+            {
+                "dp_model": self._model.serialize(),
+            }
+        )
+        return dd
+
+    @classmethod
+    def get_modifier(cls, modifier_params: dict) -> "DPModifier":
+        """Get the modifier by the parameters.
+
+        By default, all the parameters are directly passed to the constructor.
+        If not, override this method.
+
+        Parameters
+        ----------
+        modifier_params : dict
+            The modifier parameters
+
+        Returns
+        -------
+        BaseModifier
+            The modifier
+        """
+        modifier_params = modifier_params.copy()
+        modifier_params.pop("type", None)
+        # convert model_name str
+        model_name = modifier_params.pop("model_name", None)
+        if model_name is not None:
+            modifier_params["dp_model"] = None
+            modifier_params["dp_model_file_name"] = model_name
+        modifier = cls(**modifier_params)
+        return modifier

source/tests/pt/modifier/test_data_modifier.py

@@ -42,11 +42,8 @@
 from deepmd.pt.modifier.base_modifier import (
     BaseModifier,
 )
-from deepmd.pt.utils import (
-    env,
-)
-from deepmd.pt.utils.serialization import (
-    serialize_from_file,
+from deepmd.pt.modifier.dp_modifier import (
+    DPModifier,
 )
 from deepmd.pt.utils.utils import (
     to_numpy_array,
@@ -188,44 +185,19 @@ def modify_data(self, data: dict[str, Array | float], data_sys: DeepmdData) -> N
 
 
 @BaseModifier.register("scaling_tester")
-class ModifierScalingTester(BaseModifier):
-    def __new__(
-        cls,
-        *args: tuple,
-        model: str | None = None,
-        model_name: str | None = None,
-        **kwargs: dict,
-    ) -> "ModifierScalingTester":
-        return super().__new__(cls, model_name if model_name is not None else model)
-
+class ModifierScalingTester(DPModifier):
     def __init__(
         self,
-        model: torch.nn.Module | None = None,
-        model_name: str | None = None,
+        dp_model: torch.nn.Module | None = None,
+        dp_model_file_name: str | None = None,
         sfactor: float = 1.0,
         use_cache: bool = True,
     ) -> None:
         """Initialize a test modifier that applies scaled model predictions using a frozen model."""
-        super().__init__(use_cache)
+        super().__init__(dp_model, dp_model_file_name, use_cache)
         self.modifier_type = "scaling_tester"
         self.sfactor = sfactor
 
-        if model_name is None and model is None:
-            raise AttributeError("`model_name` or `model` should be specified.")
-        if model_name is not None and model is not None:
-            raise AttributeError(
-                "`model_name` and `model` cannot be used simultaneously."
-            )
-
-        if model is not None:
-            self._model = model.to(env.DEVICE)
-        if model_name is not None:
-            data = serialize_from_file(model_name)
-            self._model = EnergyModel.deserialize(data["model"]).to(env.DEVICE)
-
-        # use jit model for inference
-        self.model = torch.jit.script(self._model)
-
     def serialize(self) -> dict:
         """Serialize the modifier.
 
@@ -234,10 +206,9 @@ def serialize(self) -> dict:
         dict
             The serialized data
         """
-        dd = BaseModifier.serialize(self)
+        dd = super().serialize()
         dd.update(
             {
-                "model": self._model.serialize(),
                 "sfactor": self.sfactor,
             }
         )
@@ -249,8 +220,8 @@ def deserialize(cls, data: dict) -> "ModifierScalingTester":
         data.pop("@class", None)
         data.pop("type", None)
         data.pop("@version", None)
-        model_obj = EnergyModel.deserialize(data.pop("model"))
-        data["model"] = model_obj
+        model_obj = EnergyModel.deserialize(data.pop("dp_model"))
+        data["dp_model"] = model_obj
         return cls(**data)
 
     def forward(