Merge remote-tracking branch 'origin/devel' into descriptor-type

Signed-off-by: Jinzhe Zeng <jinzhe.zeng@rutgers.edu>

deepmd/backend/backend.py

@@ -141,6 +141,8 @@ class Feature(Flag):
         """Support Deep Eval backend."""
         NEIGHBOR_STAT = auto()
         """Support neighbor statistics."""
+        IO = auto()
+        """Support IO hook."""
 
     name: ClassVar[str] = "Unknown"
     """The formal name of the backend.
@@ -199,3 +201,27 @@ def neighbor_stat(self) -> Type["NeighborStat"]:
             The neighbor statistics of the backend.
         """
         pass
+
+    @property
+    @abstractmethod
+    def serialize_hook(self) -> Callable[[str], dict]:
+        """The serialize hook to convert the model file to a dictionary.
+
+        Returns
+        -------
+        Callable[[str], dict]
+            The serialize hook of the backend.
+        """
+        pass
+
+    @property
+    @abstractmethod
+    def deserialize_hook(self) -> Callable[[str, dict], None]:
+        """The deserialize hook to convert the dictionary to a model file.
+
+        Returns
+        -------
+        Callable[[str, dict], None]
+            The deserialize hook of the backend.
+        """
+        pass

deepmd/backend/dpmodel.py

@@ -33,7 +33,9 @@ class DPModelBackend(Backend):
 
     name = "DPModel"
     """The formal name of the backend."""
-    features: ClassVar[Backend.Feature] = Backend.Feature.NEIGHBOR_STAT
+    features: ClassVar[Backend.Feature] = (
+        Backend.Feature.NEIGHBOR_STAT | Backend.Feature.IO
+    )
     """The features of the backend."""
     suffixes: ClassVar[List[str]] = [".dp"]
     """The suffixes of the backend."""
@@ -84,3 +86,33 @@ def neighbor_stat(self) -> Type["NeighborStat"]:
         )
 
         return NeighborStat
+
+    @property
+    def serialize_hook(self) -> Callable[[str], dict]:
+        """The serialize hook to convert the model file to a dictionary.
+
+        Returns
+        -------
+        Callable[[str], dict]
+            The serialize hook of the backend.
+        """
+        from deepmd.dpmodel.utils.network import (
+            load_dp_model,
+        )
+
+        return load_dp_model
+
+    @property
+    def deserialize_hook(self) -> Callable[[str, dict], None]:
+        """The deserialize hook to convert the dictionary to a model file.
+
+        Returns
+        -------
+        Callable[[str, dict], None]
+            The deserialize hook of the backend.
+        """
+        from deepmd.dpmodel.utils.network import (
+            save_dp_model,
+        )
+
+        return save_dp_model

deepmd/backend/pytorch.py

@@ -38,6 +38,7 @@ class TensorFlowBackend(Backend):
         Backend.Feature.ENTRY_POINT
         | Backend.Feature.DEEP_EVAL
         | Backend.Feature.NEIGHBOR_STAT
+        | Backend.Feature.IO
     )
     """The features of the backend."""
     suffixes: ClassVar[List[str]] = [".pth", ".pt"]
@@ -93,3 +94,33 @@ def neighbor_stat(self) -> Type["NeighborStat"]:
         )
 
         return NeighborStat
+
+    @property
+    def serialize_hook(self) -> Callable[[str], dict]:
+        """The serialize hook to convert the model file to a dictionary.
+
+        Returns
+        -------
+        Callable[[str], dict]
+            The serialize hook of the backend.
+        """
+        from deepmd.pt.utils.serialization import (
+            serialize_from_file,
+        )
+
+        return serialize_from_file
+
+    @property
+    def deserialize_hook(self) -> Callable[[str, dict], None]:
+        """The deserialize hook to convert the dictionary to a model file.
+
+        Returns
+        -------
+        Callable[[str, dict], None]
+            The deserialize hook of the backend.
+        """
+        from deepmd.pt.utils.serialization import (
+            deserialize_to_file,
+        )
+
+        return deserialize_to_file

deepmd/backend/tensorflow.py

@@ -38,6 +38,7 @@ class TensorFlowBackend(Backend):
         Backend.Feature.ENTRY_POINT
         | Backend.Feature.DEEP_EVAL
         | Backend.Feature.NEIGHBOR_STAT
+        | Backend.Feature.IO
     )
     """The features of the backend."""
     suffixes: ClassVar[List[str]] = [".pb"]
@@ -102,3 +103,33 @@ def neighbor_stat(self) -> Type["NeighborStat"]:
         )
 
         return NeighborStat
+
+    @property
+    def serialize_hook(self) -> Callable[[str], dict]:
+        """The serialize hook to convert the model file to a dictionary.
+
+        Returns
+        -------
+        Callable[[str], dict]
+            The serialize hook of the backend.
+        """
+        from deepmd.tf.utils.serialization import (
+            serialize_from_file,
+        )
+
+        return serialize_from_file
+
+    @property
+    def deserialize_hook(self) -> Callable[[str, dict], None]:
+        """The deserialize hook to convert the dictionary to a model file.
+
+        Returns
+        -------
+        Callable[[str, dict], None]
+            The deserialize hook of the backend.
+        """
+        from deepmd.tf.utils.serialization import (
+            deserialize_to_file,
+        )
+
+        return deserialize_to_file

deepmd/dpmodel/utils/network.py

@@ -6,6 +6,9 @@
 import copy
 import itertools
 import json
+from datetime import (
+    datetime,
+)
 from typing import (
     ClassVar,
     Dict,
@@ -54,6 +57,8 @@ def traverse_model_dict(model_obj, callback: callable, is_variable: bool = False
     elif isinstance(model_obj, list):
         for ii, vv in enumerate(model_obj):
             model_obj[ii] = traverse_model_dict(vv, callback, is_variable=is_variable)
+    elif model_obj is None:
+        return model_obj
     elif is_variable:
         model_obj = callback(model_obj)
     return model_obj
@@ -79,7 +84,8 @@ def __call__(self):
         return self.count
 
 
-def save_dp_model(filename: str, model_dict: dict, extra_info: Optional[dict] = None):
+# TODO: should be moved to otherwhere...
+def save_dp_model(filename: str, model_dict: dict) -> None:
     """Save a DP model to a file in the native format.
 
     Parameters
@@ -88,15 +94,9 @@ def save_dp_model(filename: str, model_dict: dict, extra_info: Optional[dict] =
         The filename to save to.
     model_dict : dict
         The model dict to save.
-    extra_info : dict, optional
-        Extra meta information to save.
     """
     model_dict = model_dict.copy()
     variable_counter = Counter()
-    if extra_info is not None:
-        extra_info = extra_info.copy()
-    else:
-        extra_info = {}
     with h5py.File(filename, "w") as f:
         model_dict = traverse_model_dict(
             model_dict,
@@ -105,10 +105,11 @@ def save_dp_model(filename: str, model_dict: dict, extra_info: Optional[dict] =
             ).name,
         )
         save_dict = {
-            "model": model_dict,
             "software": "deepmd-kit",
             "version": __version__,
-            **extra_info,
+            # use UTC+0 time
+            "time": str(datetime.utcnow()),
+            **model_dict,
         }
         f.attrs["json"] = json.dumps(save_dict, separators=(",", ":"))
 

deepmd/entrypoints/convert_backend.py

@@ -0,0 +1,27 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+from deepmd.backend.backend import (
+    Backend,
+)
+
+
+def convert_backend(
+    *,  # Enforce keyword-only arguments
+    INPUT: str,
+    OUTPUT: str,
+    **kwargs,
+) -> None:
+    """Convert a model file from one backend to another.
+
+    Parameters
+    ----------
+    INPUT : str
+        The input model file.
+    INPUT : str
+        The output model file.
+    """
+    inp_backend: Backend = Backend.detect_backend_by_model(INPUT)()
+    out_backend: Backend = Backend.detect_backend_by_model(OUTPUT)()
+    inp_hook = inp_backend.serialize_hook
+    out_hook = out_backend.deserialize_hook
+    data = inp_hook(INPUT)
+    out_hook(OUTPUT, data)

deepmd/entrypoints/main.py

@@ -12,6 +12,9 @@
 from deepmd.backend.suffix import (
     format_model_suffix,
 )
+from deepmd.entrypoints.convert_backend import (
+    convert_backend,
+)
 from deepmd.entrypoints.doc import (
     doc_train_input,
 )
@@ -76,5 +79,7 @@ def main(args: argparse.Namespace):
         neighbor_stat(**dict_args)
     elif args.command == "gui":
         start_dpgui(**dict_args)
+    elif args.command == "convert-backend":
+        convert_backend(**dict_args)
     else:
         raise ValueError(f"Unknown command: {args.command}")

deepmd/main.py

@@ -721,6 +721,23 @@ def main_parser() -> argparse.ArgumentParser:
             "to the network on both IPv4 and IPv6 (where available)."
         ),
     )
+
+    # convert_backend
+    parser_convert_backend = subparsers.add_parser(
+        "convert-backend",
+        parents=[parser_log],
+        help="Convert model to another backend.",
+        formatter_class=RawTextArgumentDefaultsHelpFormatter,
+        epilog=textwrap.dedent(
+            """\
+        examples:
+            dp convert-backend model.pb model.pth
+            dp convert-backend model.pb model.dp
+        """
+        ),
+    )
+    parser_convert_backend.add_argument("INPUT", help="The input model file.")
+    parser_convert_backend.add_argument("OUTPUT", help="The output model file.")
     return parser
 
 
@@ -767,6 +784,7 @@ def main():
         "model-devi",
         "neighbor-stat",
         "gui",
+        "convert-backend",
     ):
         # common entrypoints
         from deepmd.entrypoints.main import main as deepmd_main

deepmd/pt/model/atomic_model/base_atomic_model.py

@@ -14,3 +14,7 @@ class BaseAtomicModel(BaseAtomicModel_):
     # export public methods that are not abstract
     get_nsel = torch.jit.export(BaseAtomicModel_.get_nsel)
     get_nnei = torch.jit.export(BaseAtomicModel_.get_nnei)
+
+    @torch.jit.export
+    def get_model_def_script(self) -> str:
+        return self.model_def_script

deepmd/pt/model/atomic_model/dp_atomic_model.py

@@ -50,6 +50,7 @@ class DPAtomicModel(torch.nn.Module, BaseAtomicModel):
 
     def __init__(self, descriptor, fitting, type_map: Optional[List[str]]):
         super().__init__()
+        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

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

deepmd/pt/model/atomic_model/pairtab_atomic_model.py

@@ -52,6 +52,7 @@ def __init__(
         self, tab_file: str, rcut: float, sel: Union[int, List[int]], **kwargs
     ):
         super().__init__()
+        self.model_def_script = ""
         self.tab_file = tab_file
         self.rcut = rcut
         self.tab = self._set_pairtab(tab_file, rcut)

deepmd/pt/model/descriptor/se_atten.py

@@ -87,6 +87,8 @@ def __init__(
         self.ffn = ffn
         self.ffn_embed_dim = ffn_embed_dim
         self.activation = activation
+        # TODO: To be fixed: precision should be given from inputs
+        self.prec = torch.float64
         self.scaling_factor = scaling_factor
         self.head_num = head_num
         self.normalize = normalize

deepmd/pt/model/model/__init__.py

@@ -10,6 +10,7 @@
 """
 
 import copy
+import json
 
 from deepmd.pt.model.atomic_model import (
     DPAtomicModel,
@@ -98,7 +99,9 @@ def get_model(model_params):
             fitting_net["return_energy"] = True
     fitting = Fitting(**fitting_net)
 
-    return EnergyModel(descriptor, fitting, type_map=model_params["type_map"])
+    model = EnergyModel(descriptor, fitting, type_map=model_params["type_map"])
+    model.model_def_script = json.dumps(model_params)
+    return model
 
 
 __all__ = [