feat(pt): consistent "frozen" model

deepmd/dpmodel/atomic_model/make_base_atomic_model.py

@@ -178,10 +178,6 @@ def do_grad_(self, var_name: str, base: str) -> bool:
                 return self.fitting_output_def()[var_name].c_differentiable
             return self.fitting_output_def()[var_name].r_differentiable
 
-        def get_model_def_script(self) -> str:
-            # TODO: implement this method; saved to model
-            raise NotImplementedError
-
     setattr(BAM, fwd_method_name, BAM.fwd)
     delattr(BAM, "fwd")
 

deepmd/dpmodel/model/base_model.py

@@ -172,3 +172,10 @@
     deepmd.dpmodel.model.base_model.BaseBaseModel
         Backend-independent BaseModel class.
     """
+
+    def __init__(self) -> None:
+        self.model_def_script = ""
+
+    def get_model_def_script(self) -> str:
+        """Get the model definition script."""
+        return self.model_def_script

deepmd/dpmodel/model/make_model.py

@@ -73,6 +73,7 @@
             atomic_model_: Optional[T_AtomicModel] = None,
             **kwargs,
         ):
+            BaseModel.__init__(self)
             if atomic_model_ is not None:
                 self.atomic_model: T_AtomicModel = atomic_model_
             else:
@@ -452,10 +453,6 @@
             """Returns the total number of selected neighboring atoms in the cut-off radius."""
             return self.atomic_model.get_nnei()
 
-        def get_model_def_script(self) -> str:
-            """Get the model definition script."""
-            return self.atomic_model.get_model_def_script()
-
         def get_sel(self) -> List[int]:
             """Returns the number of selected atoms for each type."""
             return self.atomic_model.get_sel()

deepmd/dpmodel/utils/network.py

@@ -230,6 +230,10 @@
             variables.get("b", None),
             variables.get("idt", None),
         )
+        if obj.b is not None:
+            obj.b = obj.b.ravel()
+        if obj.idt is not None:
+            obj.idt = obj.idt.ravel()
         obj.check_shape_consistency()
         return obj
 

deepmd/pt/model/atomic_model/base_atomic_model.py

@@ -58,9 +58,6 @@ def reinit_pair_exclude(
         else:
             self.pair_excl = PairExcludeMask(self.get_ntypes(), self.pair_exclude_types)
 
-    def get_model_def_script(self) -> str:
-        return self.model_def_script
-
     def atomic_output_def(self) -> FittingOutputDef:
         old_def = self.fitting_output_def()
         if self.atom_excl is None:

deepmd/pt/model/atomic_model/dp_atomic_model.py

@@ -55,7 +55,6 @@ def __init__(
         **kwargs,
     ):
         torch.nn.Module.__init__(self)
-        self.model_def_script = ""
         ntypes = len(type_map)
         self.type_map = type_map
         self.ntypes = ntypes

deepmd/pt/model/atomic_model/linear_atomic_model.py

@@ -360,7 +360,6 @@ def __init__(
     ):
         models = [dp_model, zbl_model]
         super().__init__(models, type_map, **kwargs)
-        self.model_def_script = ""
         self.dp_model = dp_model
         self.zbl_model = zbl_model
 

deepmd/pt/model/atomic_model/pairtab_atomic_model.py

@@ -78,7 +78,6 @@ def __init__(
         **kwargs,
     ):
         torch.nn.Module.__init__(self)
-        self.model_def_script = ""
         self.tab_file = tab_file
         self.rcut = rcut
         self.tab = self._set_pairtab(tab_file, rcut)

deepmd/pt/model/model/__init__.py

@@ -37,6 +37,9 @@
 from .ener_model import (
     EnergyModel,
 )
+from .frozen import (
+    FrozenModel,
+)
 from .make_hessian_model import (
     make_hessian_model,
 )
@@ -173,6 +176,7 @@
     "get_model",
     "DPModel",
     "EnergyModel",
+    "FrozenModel",
     "SpinModel",
     "SpinEnergyModel",
     "DPZBLModel",

deepmd/pt/model/model/frozen.py

@@ -0,0 +1,174 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+import json
+import tempfile
+from typing import (
+    Dict,
+    List,
+    Optional,
+)
+
+import torch
+
+from deepmd.dpmodel.output_def import (
+    FittingOutputDef,
+)
+from deepmd.entrypoints.convert_backend import (
+    convert_backend,
+)
+from deepmd.pt.model.model.model import (
+    BaseModel,
+)
+
+
+@BaseModel.register("frozen")
+class FrozenModel(BaseModel):
+    """Load model from a frozen model, which cannot be trained.
+
+    Parameters
+    ----------
+    model_file : str
+        The path to the frozen model
+    """
+
+    def __init__(self, model_file: str, **kwargs):
+        super().__init__(**kwargs)
+        self.model_file = model_file
+        if model_file.endswith(".pth"):
+            self.model = torch.jit.load(model_file)
+        else:
+            # try to convert from other formats
+            with tempfile.NamedTemporaryFile(suffix=".pth") as f:
+                convert_backend(INPUT=model_file, OUTPUT=f.name)
+                self.model = torch.jit.load(f.name)
+
+    @torch.jit.export
+    def fitting_output_def(self) -> FittingOutputDef:
+        """Get the output def of developer implemented atomic models."""
+        return self.model.fitting_output_def()
+
+    @torch.jit.export
+    def get_rcut(self) -> float:
+        """Get the cut-off radius."""
+        return self.model.get_rcut()
+
+    @torch.jit.export
+    def get_type_map(self) -> List[str]:
+        """Get the type map."""
+        return self.model.get_type_map()
+
+    @torch.jit.export
+    def get_sel(self) -> List[int]:
+        """Returns the number of selected atoms for each type."""
+        return self.model.get_sel()
+
+    @torch.jit.export
+    def get_dim_fparam(self) -> int:
+        """Get the number (dimension) of frame parameters of this atomic model."""
+        return self.model.get_dim_fparam()
+
+    @torch.jit.export
+    def get_dim_aparam(self) -> int:
+        """Get the number (dimension) of atomic parameters of this atomic model."""
+        return self.model.get_dim_aparam()
+
+    @torch.jit.export
+    def get_sel_type(self) -> List[int]:
+        """Get the selected atom types of this model.
+
+        Only atoms with selected atom types have atomic contribution
+        to the result of the model.
+        If returning an empty list, all atom types are selected.
+        """
+        return self.model.get_sel_type()
+
+    @torch.jit.export
+    def is_aparam_nall(self) -> bool:
+        """Check whether the shape of atomic parameters is (nframes, nall, ndim).
+
+        If False, the shape is (nframes, nloc, ndim).
+        """
+        return self.model.is_aparam_nall()
+
+    @torch.jit.export
+    def mixed_types(self) -> bool:
+        """If true, the model
+        1. assumes total number of atoms aligned across frames;
+        2. uses a neighbor list that does not distinguish different atomic types.
+
+        If false, the model
+        1. assumes total number of atoms of each atom type aligned across frames;
+        2. uses a neighbor list that distinguishes different atomic types.
+
+        """
+        return self.model.mixed_types()
+
+    @torch.jit.export
+    def forward(
+        self,
+        coord,
+        atype,
+        box: Optional[torch.Tensor] = None,
+        fparam: Optional[torch.Tensor] = None,
+        aparam: Optional[torch.Tensor] = None,
+        do_atomic_virial: bool = False,
+    ) -> Dict[str, torch.Tensor]:
+        return self.model.forward(
+            coord,
+            atype,
+            box=box,
+            fparam=fparam,
+            aparam=aparam,
+            do_atomic_virial=do_atomic_virial,
+        )
+
+    @torch.jit.export
+    def get_model_def_script(self) -> str:
+        """Get the model definition script."""
+        # try to use the original script instead of "frozen model"
+        # Note: this cannot change the script of the parent model
+        # it may still try to load hard-coded filename, which might
+        # be a problem
+        return self.model.get_model_def_script()
+
+    def serialize(self) -> dict:
+        from deepmd.pt.model.model import (
+            get_model,
+        )
+
+        # try to recover the original model
+        model_def_script = json.loads(self.get_model_def_script())
+        model = get_model(model_def_script)
+        model.load_state_dict(self.model.state_dict())
+        return model.serialize()
+
+    @classmethod
+    def deserialize(cls, data: dict):
+        raise RuntimeError("Should not touch here.")
+
+    @torch.jit.export
+    def get_nnei(self) -> int:
+        """Returns the total number of selected neighboring atoms in the cut-off radius."""
+        return self.model.get_nnei()
+
+    @torch.jit.export
+    def get_nsel(self) -> int:
+        """Returns the total number of selected neighboring atoms in the cut-off radius."""
+        return self.model.get_nsel()
+
+    @classmethod
+    def update_sel(cls, global_jdata: dict, local_jdata: dict):
+        """Update the selection and perform neighbor statistics.
+
+        Parameters
+        ----------
+        global_jdata : dict
+            The global data, containing the training section
+        local_jdata : dict
+            The local data refer to the current class
+        """
+        return local_jdata
+
+    @torch.jit.export
+    def model_output_type(self) -> str:
+        """Get the output type for the model."""
+        return self.model.model_output_type()

deepmd/pt/model/model/make_model.py

@@ -469,11 +469,6 @@ def get_nnei(self) -> int:
             """Returns the total number of selected neighboring atoms in the cut-off radius."""
             return self.atomic_model.get_nnei()
 
-        @torch.jit.export
-        def get_model_def_script(self) -> str:
-            """Get the model definition script."""
-            return self.atomic_model.get_model_def_script()
-
         def atomic_output_def(self) -> FittingOutputDef:
             """Get the output def of the atomic model."""
             return self.atomic_model.atomic_output_def()

deepmd/pt/model/model/model.py

@@ -17,6 +17,7 @@
     def __init__(self, *args, **kwargs):
         """Construct a basic model for different tasks."""
         torch.nn.Module.__init__(self)
+        self.model_def_script = ""
 
     def compute_or_load_stat(
         self,
@@ -39,3 +40,8 @@
             The path to the statistics files.
         """
         raise NotImplementedError
+
+    @torch.jit.export
+    def get_model_def_script(self) -> str:
+        """Get the model definition script."""
+        return self.model_def_script

deepmd/tf/fit/ener.py

@@ -213,7 +213,7 @@
             else:
                 self.atom_ener.append(None)
         self.useBN = False
-        self.bias_atom_e = np.zeros(self.ntypes, dtype=np.float64)
+        self.bias_atom_e = np.zeros((self.ntypes, 1), dtype=np.float64)
         # data requirement
         if self.numb_fparam > 0:
             add_data_requirement(

deepmd/tf/model/frozen.py

@@ -1,4 +1,7 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
+import json
+import os
+import tempfile
 from enum import (
     Enum,
 )
@@ -7,6 +10,9 @@
     Union,
 )
 
+from deepmd.entrypoints.convert_backend import (
+    convert_backend,
+)
 from deepmd.infer.deep_pot import (
     DeepPot,
 )
@@ -24,6 +30,10 @@
 from deepmd.tf.loss.loss import (
     Loss,
 )
+from deepmd.tf.utils.graph import (
+    get_tensor_by_name_from_graph,
+    load_graph_def,
+)
 
 from .model import (
     Model,
@@ -43,7 +53,14 @@
     def __init__(self, model_file: str, **kwargs):
         super().__init__(**kwargs)
         self.model_file = model_file
-        self.model = DeepPotential(model_file)
+        if not model_file.endswith(".pb"):
+            # try to convert from other formats
+            with tempfile.NamedTemporaryFile(
+                suffix=".pb", dir=os.curdir, delete=False
+            ) as f:
+                convert_backend(INPUT=model_file, OUTPUT=f.name)
+                self.model_file = f.name
+        self.model = DeepPotential(self.model_file)
         if isinstance(self.model, DeepPot):
             self.model_type = "ener"
         else:
@@ -228,3 +245,19 @@
         """
         # we don't know how to compress it, so no neighbor statistics here
         return local_jdata
+
+    def serialize(self, suffix: str = "") -> dict:
+        # try to recover the original model
+        # the current graph contains a prefix "load",
+        # so it cannot used to recover the original model
+        graph, graph_def = load_graph_def(self.model_file)
+        t_jdata = get_tensor_by_name_from_graph(graph, "train_attr/training_script")
+        jdata = json.loads(t_jdata)
+        model = Model(**jdata["model"])
+        # important! must be called before serialize
+        model.init_variables(graph=graph, graph_def=graph_def)
+        return model.serialize()
+
+    @classmethod
+    def deserialize(cls, data: dict, suffix: str = ""):
+        raise RuntimeError("Should not touch here.")

deepmd/tf/model/model.py

@@ -566,7 +566,8 @@ def deserialize(cls, data: dict, suffix: str = "") -> "Model":
         """
         if cls is Model:
             return Model.get_class_by_type(data.get("type", "standard")).deserialize(
-                data
+                data,
+                suffix=suffix,
             )
         raise NotImplementedError("Not implemented in class %s" % cls.__name__)
 

deepmd/utils/argcheck.py

@@ -1461,7 +1461,6 @@ def frozen_model_args() -> Argument:
         [
             Argument("model_file", str, optional=False, doc=doc_model_file),
         ],
-        doc=doc_only_tf_supported,
     )
     return ca