initial commit with se_e2_a supported only

Signed-off-by: Jinzhe Zeng <[email protected]>

deepmd/descriptor/descriptor.py

@@ -509,3 +509,31 @@ def update_sel(cls, global_jdata: dict, local_jdata: dict):
         # call subprocess
         cls = cls.get_class_by_input(local_jdata)
         return cls.update_sel(global_jdata, local_jdata)
+
+    @classmethod
+    def deserialize(cls, data: dict):
+        """Deserialize the model.
+
+        Parameters
+        ----------
+        data : dict
+            The serialized data
+
+        Returns
+        -------
+        Model
+            The deserialized model
+        """
+        if cls is Descriptor:
+            return Descriptor.get_class_by_input(data).deserialize(data)
+        raise NotImplementedError("Not implemented in class %s" % cls.__name__)
+
+    def serialize(self) -> dict:
+        """Serialize the model.
+
+        Returns
+        -------
+        dict
+            The serialized data
+        """
+        raise NotImplementedError("Not implemented in class %s" % self.__name__)

deepmd/descriptor/se_a.py

@@ -192,6 +192,7 @@ def __init__(
         self.seed_shift = embedding_net_rand_seed_shift(self.filter_neuron)
         self.trainable = trainable
         self.compress_activation_fn = get_activation_func(activation_function)
+        self.activation_function_name = activation_function
         self.filter_activation_fn = get_activation_func(activation_function)
         self.filter_precision = get_precision(precision)
         self.filter_np_precision = get_np_precision(precision)
@@ -1334,3 +1335,57 @@ def explicit_ntypes(self) -> bool:
         if self.stripped_type_embedding:
             return True
         return False
+
+    @classmethod
+    def deserialize(cls, data: dict):
+        """Deserialize the model.
+
+        Parameters
+        ----------
+        data : dict
+            The serialized data
+
+        Returns
+        -------
+        Model
+            The deserialized model
+        """
+        descriptor = cls(**data)
+        descriptor.davg = data["@variables"]["davg"]
+        descriptor.dstd = data["@variables"]["dstd"]
+        descriptor.embedding_net_variables = data["@variables"]
+        descriptor.original_sel = data["@variables"]["original_sel"]
+        return descriptor
+
+    def serialize(self) -> dict:
+        """Serialize the model.
+
+        Returns
+        -------
+        dict
+            The serialized data
+        """
+        return {
+            "type": "se_e2_a",
+            "rcut": self.rcut_r,
+            "rcut_smth": self.rcut_r_smth,
+            "sel": self.sel_a,
+            "neuron": self.filter_neuron,
+            "axis_neuron": self.n_axis_neuron,
+            "resnet_dt": self.filter_resnet_dt,
+            "trainable": self.trainable,
+            "seed": self.seed,
+            "type_one_side": self.type_one_side,
+            "exclude_types": list(self.exclude_types),
+            "set_davg_zero": self.set_davg_zero,
+            "activation_function": self.activation_function_name,
+            "precision": self.filter_precision.name,
+            "uniform_seed": self.uniform_seed,
+            "stripped_type_embedding": self.stripped_type_embedding,
+            "@variables": {
+                **self.embedding_net_variables,
+                "davg": self.davg,
+                "dstd": self.dstd,
+                "original_sel": self.original_sel,
+            },
+        }

deepmd/entrypoints/convert.py

@@ -9,6 +9,10 @@
     convert_pbtxt_to_pb,
     convert_to_21,
 )
+from deepmd.utils.convert_dp import (
+    convert_dp_to_pb,
+    convert_pb_to_dp,
+)
 
 
 def convert(
@@ -23,6 +27,11 @@ def convert(
             convert_pb_to_pbtxt(input_model, output_model)
         else:
             raise RuntimeError("input model is already pbtxt")
+    elif output_model.endswith(".dp"):
+        if input_model.endswith(".pb"):
+            convert_pb_to_dp(input_model, output_model)
+        else:
+            raise RuntimeError("Unsupported format")
     else:
         if FROM == "auto":
             convert_to_21(input_model, output_model)
@@ -39,5 +48,7 @@ def convert(
             convert_20_to_21(input_model, output_model)
         elif FROM == "pbtxt":
             convert_pbtxt_to_pb(input_model, output_model)
+        elif FROM == "dp":
+            convert_dp_to_pb(input_model, output_model)
         else:
             raise RuntimeError("unsupported model version " + FROM)

deepmd/fit/ener.py

@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
 import logging
 from typing import (
+    TYPE_CHECKING,
     List,
     Optional,
 )
@@ -53,6 +54,11 @@
     Spin,
 )
 
+if TYPE_CHECKING:
+    from deepmd.descriptor import (
+        Descriptor,
+    )
+
 log = logging.getLogger(__name__)
 
 
@@ -130,7 +136,7 @@ class EnerFitting(Fitting):
 
     def __init__(
         self,
-        descrpt: tf.Tensor,
+        descrpt: "Descriptor",
         neuron: List[int] = [120, 120, 120],
         resnet_dt: bool = True,
         numb_fparam: int = 0,
@@ -176,6 +182,7 @@ def __init__(
         self.ntypes_spin = self.spin.get_ntypes_spin() if self.spin is not None else 0
         self.seed_shift = one_layer_rand_seed_shift()
         self.tot_ener_zero = tot_ener_zero
+        self.activation_function_name = activation_function
         self.fitting_activation_fn = get_activation_func(activation_function)
         self.fitting_precision = get_precision(precision)
         self.trainable = trainable
@@ -916,3 +923,66 @@ def get_loss(self, loss: dict, lr) -> Loss:
             return EnerSpinLoss(**loss, use_spin=self.spin.use_spin)
         else:
             raise RuntimeError("unknown loss type")
+
+    @classmethod
+    def deserialize(cls, data: dict):
+        """Deserialize the model.
+
+        Parameters
+        ----------
+        data : dict
+            The serialized data
+
+        Returns
+        -------
+        Model
+            The deserialized model
+        """
+        fitting = cls(**data)
+        fitting.fitting_net_variables = data["@variables"]
+        fitting.bias_atom_e = fitting.fitting_net_variables.pop("bias_atom_e")
+        if fitting.numb_fparam > 0:
+            fitting.fparam_avg = fitting.fitting_net_variables.pop("fparam_avg")
+            fitting.fparam_inv_std = fitting.fitting_net_variables.pop("fparam_inv_std")
+        if fitting.numb_aparam > 0:
+            fitting.aparam_avg = fitting.fitting_net_variables.pop("aparam_avg")
+            fitting.aparam_inv_std = fitting.fitting_net_variables.pop("aparam_inv_std")
+        return fitting
+
+    def serialize(self) -> dict:
+        """Serialize the model.
+
+        Returns
+        -------
+        dict
+            The serialized data
+        """
+        data = {
+            "type": "ener",
+            "neuron": self.n_neuron,
+            "resnet_dt": self.resnet_dt,
+            "numb_fparam": self.numb_fparam,
+            "numb_aparam": self.numb_aparam,
+            "rcond": self.rcond,
+            "tot_ener_zero": self.tot_ener_zero,
+            "trainable": self.trainable,
+            "seed": self.seed,
+            "atom_ener": self.atom_ener,
+            "activation_function": self.activation_function_name,
+            "precision": self.fitting_precision.name,
+            "uniform_seed": self.uniform_seed,
+            "layer_name": self.layer_name,
+            "use_aparam_as_mask": self.use_aparam_as_mask,
+            "@variables": {
+                **self.fitting_net_variables,
+                "bias_atom_e": self.bias_atom_e,
+            },
+        }
+
+        if self.numb_fparam > 0:
+            data["@variables"]["fparam_avg"] = self.fparam_avg
+            data["@variables"]["fparam_inv_std"] = self.fparam_inv_std
+        if self.numb_aparam > 0:
+            data["@variables"]["aparam_avg"] = self.aparam_avg
+            data["@variables"]["aparam_inv_std"] = self.aparam_inv_std
+        return data

deepmd/fit/fitting.py

@@ -43,16 +43,20 @@ class SomeFitting(Fitting):
         """
         return Fitting.__plugins.register(key)
 
+    @classmethod
+    def get_class_by_input(cls, input: dict):
+        try:
+            fitting_type = input["type"]
+        except KeyError:
+            raise KeyError("the type of fitting should be set by `type`")
+        if fitting_type in Fitting.__plugins.plugins:
+            return Fitting.__plugins.plugins[fitting_type]
+        else:
+            raise RuntimeError("Unknown descriptor type: " + fitting_type)
+
     def __new__(cls, *args, **kwargs):
         if cls is Fitting:
-            try:
-                fitting_type = kwargs["type"]
-            except KeyError:
-                raise KeyError("the type of fitting should be set by `type`")
-            if fitting_type in Fitting.__plugins.plugins:
-                cls = Fitting.__plugins.plugins[fitting_type]
-            else:
-                raise RuntimeError("Unknown descriptor type: " + fitting_type)
+            cls = Fitting.get_class_by_input(kwargs)
         return super().__new__(cls)
 
     @property
@@ -102,3 +106,31 @@ def get_loss(self, loss: dict, lr) -> Loss:
         Loss
             the loss function
         """
+
+    @classmethod
+    def deserialize(cls, data: dict):
+        """Deserialize the model.
+
+        Parameters
+        ----------
+        data : dict
+            The serialized data
+
+        Returns
+        -------
+        Model
+            The deserialized model
+        """
+        if cls is Fitting:
+            return Fitting.get_class_by_input(data).deserialize(data)
+        raise NotImplementedError("Not implemented in class %s" % cls.__name__)
+
+    def serialize(self) -> dict:
+        """Serialize the model.
+
+        Returns
+        -------
+        dict
+            The serialized data
+        """
+        raise NotImplementedError("Not implemented in class %s" % self.__name__)

deepmd/model/ener.py

@@ -7,12 +7,18 @@
 
 import numpy as np
 
+from deepmd.descriptor.descriptor import (
+    Descriptor,
+)
 from deepmd.env import (
     MODEL_VERSION,
     global_cvt_2_ener_float,
     op_module,
     tf,
 )
+from deepmd.fit.fitting import (
+    Fitting,
+)
 from deepmd.utils.data_system import (
     DeepmdDataSystem,
 )
@@ -512,3 +518,59 @@ def change_energy_bias(
             bias_shift,
             self.data_bias_nsample,
         )
+
+    @classmethod
+    def deserialize(cls, data: dict):
+        """Deserialize the model.
+
+        Parameters
+        ----------
+        data : dict
+            The serialized data
+
+        Returns
+        -------
+        Model
+            The deserialized model
+        """
+        data = data.copy()
+        data["descriptor"] = Descriptor.deserialize(data["descriptor"])
+        fitting_data = data["fitting_net"].copy()
+        fitting_data["descrpt"] = data["descriptor"]
+        data["fitting_net"] = Fitting.deserialize(fitting_data)
+        if data.get("type_embedding") is not None:
+            data["type_embedding"] = TypeEmbedNet.deserialize(data["type_embedding"])
+        return cls(**data)
+
+    def serialize(self) -> dict:
+        """Serialize the model.
+
+        Returns
+        -------
+        dict
+            The serialized data
+        """
+        data = {
+            "type": "standard",
+            "descriptor": self.descrpt.serialize(),
+            "fitting_net": self.fitting.serialize(),
+            "type_embedding": self.typeebd.serialize()
+            if self.typeebd is not None
+            else None,
+            "type_map": self.get_type_map(),
+            "data_stat_nbatch": self.data_stat_nbatch,
+            "data_stat_protect": self.data_stat_protect,
+            "use_srtab": self.srtab_name,
+            "spin": self.spin.serialize() if self.spin is not None else None,
+            "data_bias_nsample": self.data_bias_nsample,
+        }
+        if self.srtab_name is not None:
+            data.merge(
+                {
+                    "srtab_add_bias": self.srtab_add_bias,
+                    "smin_alpha": self.smin_alpha,
+                    "sw_rmin": self.sw_rmin,
+                    "sw_rmax": self.sw_rmax,
+                }
+            )
+        return data