support TF se_e2_a serialization; add a common test fixture to compare TF, PT, and DP models

deepmd/dpmodel/common.py

@@ -6,6 +6,10 @@
 
 import numpy as np
 
+from deepmd.common import (
+    GLOBAL_NP_FLOAT_PRECISION,
+)
+
 PRECISION_DICT = {
     "float16": np.float16,
     "float32": np.float32,
@@ -15,6 +19,7 @@
     "double": np.float64,
     "int32": np.int32,
     "int64": np.int64,
+    "default": GLOBAL_NP_FLOAT_PRECISION,
 }
 DEFAULT_PRECISION = "float64"
 

deepmd/dpmodel/descriptor/se_e2_a.py

@@ -15,6 +15,7 @@
 
 from deepmd.dpmodel import (
     DEFAULT_PRECISION,
+    PRECISION_DICT,
     NativeOP,
 )
 from deepmd.dpmodel.utils import (
@@ -133,6 +134,8 @@
         activation_function: str = "tanh",
         precision: str = DEFAULT_PRECISION,
         spin: Optional[Any] = None,
+        # consistent with argcheck, not used though
+        seed: Optional[int] = None,
     ) -> None:
         ## seed, uniform_seed, multi_task, not included.
         if not type_one_side:
@@ -163,6 +166,8 @@
             ndim=(1 if self.type_one_side else 2),
             network_type="embedding_network",
         )
+        if not self.type_one_side:
+            raise NotImplementedError("type_one_side == False not implemented")
         for ii in range(self.ntypes):
             self.embeddings[(ii,)] = EmbeddingNet(
                 in_dim,
@@ -316,7 +321,8 @@
             "exclude_types": self.exclude_types,
             "set_davg_zero": self.set_davg_zero,
             "activation_function": self.activation_function,
-            "precision": self.precision,
+            # make deterministic
+            "precision": np.dtype(PRECISION_DICT[self.precision]).name,
             "spin": self.spin,
             "env_mat": self.env_mat.serialize(),
             "embeddings": self.embeddings.serialize(),

deepmd/dpmodel/utils/network.py

@@ -192,7 +192,8 @@ def serialize(self) -> dict:
             "use_timestep": self.idt is not None,
             "activation_function": self.activation_function,
             "resnet": self.resnet,
-            "precision": self.precision,
+            # make deterministic
+            "precision": np.dtype(PRECISION_DICT[self.precision]).name,
             "@variables": data,
         }
 
@@ -464,7 +465,8 @@ def serialize(self) -> dict:
                 "neuron": self.neuron.copy(),
                 "activation_function": self.activation_function,
                 "resnet_dt": self.resnet_dt,
-                "precision": self.precision,
+                # make deterministic
+                "precision": np.dtype(PRECISION_DICT[self.precision]).name,
                 "layers": [layer.serialize() for layer in self.layers],
             }
 

deepmd/pt/model/descriptor/se_a.py

@@ -20,6 +20,7 @@
 )
 from deepmd.pt.utils.env import (
     PRECISION_DICT,
+    RESERVED_PRECISON_DICT,
 )
 
 try:
@@ -207,7 +208,8 @@ def serialize(self) -> dict:
             "resnet_dt": obj.resnet_dt,
             "set_davg_zero": obj.set_davg_zero,
             "activation_function": obj.activation_function,
-            "precision": obj.precision,
+            # make deterministic
+            "precision": RESERVED_PRECISON_DICT[obj.prec],
             "embeddings": obj.filter_layers.serialize(),
             "env_mat": DPEnvMat(obj.rcut, obj.rcut_smth).serialize(),
             "@variables": {
@@ -223,6 +225,7 @@ def serialize(self) -> dict:
 
     @classmethod
     def deserialize(cls, data: dict) -> "DescrptSeA":
+        data = data.copy()
         variables = data.pop("@variables")
         embeddings = data.pop("embeddings")
         env_mat = data.pop("env_mat")

deepmd/pt/utils/env.py

@@ -42,6 +42,15 @@
     "int64": torch.int64,
 }
 GLOBAL_PT_FLOAT_PRECISION = PRECISION_DICT[np.dtype(GLOBAL_NP_FLOAT_PRECISION).name]
+PRECISION_DICT["default"] = GLOBAL_PT_FLOAT_PRECISION
+# cannot automatically generated
+RESERVED_PRECISON_DICT = {
+    torch.float16: "float16",
+    torch.float32: "float32",
+    torch.float64: "float64",
+    torch.int32: "int32",
+    torch.int64: "int64",
+}
 DEFAULT_PRECISION = "float64"
 
 # throw warnings if threads not set
@@ -58,6 +67,7 @@
     "GLOBAL_PT_FLOAT_PRECISION",
     "DEFAULT_PRECISION",
     "PRECISION_DICT",
+    "RESERVED_PRECISON_DICT",
     "SAMPLER_RECORD",
     "NUM_WORKERS",
     "DEVICE",

deepmd/tf/descriptor/descriptor.py

@@ -509,3 +509,41 @@
         # call subprocess
         cls = cls.get_class_by_input(local_jdata)
         return cls.update_sel(global_jdata, local_jdata)
+
+    @classmethod
+    def deserialize(cls, data: dict, suffix: str = "") -> "Descriptor":
+        """Deserialize the model.
+
+        There is no suffix in a native DP model, but it is important
+        for the TF backend.
+
+        Parameters
+        ----------
+        data : dict
+            The serialized data
+        suffix : str, optional
+            Name suffix to identify this descriptor
+
+        Returns
+        -------
+        Descriptor
+            The deserialized descriptor
+        """
+        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, suffix: str = "") -> dict:
+        """Serialize the model.
+
+        There is no suffix in a native DP model, but it is important
+        for the TF backend.
+
+        Returns
+        -------
+        dict
+            The serialized data
+        suffix : str, optional
+            Name suffix to identify this descriptor
+        """
+        raise NotImplementedError("Not implemented in class %s" % self.__name__)

deepmd/tf/descriptor/se.py

@@ -1,9 +1,17 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
+import re
 from typing import (
+    List,
+    Set,
     Tuple,
 )
 
+from deepmd.dpmodel.utils.network import (
+    EmbeddingNet,
+    NetworkCollection,
+)
 from deepmd.tf.env import (
+    EMBEDDING_NET_PATTERN,
     tf,
 )
 from deepmd.tf.utils.graph import (
@@ -160,3 +168,166 @@
         # default behavior is to update sel which is a list
         local_jdata_cpy = local_jdata.copy()
         return update_one_sel(global_jdata, local_jdata_cpy, False)
+
+    def serialize_network(
+        self,
+        ntypes: int,
+        ndim: int,
+        in_dim: int,
+        neuron: List[int],
+        activation_function: str,
+        resnet_dt: bool,
+        variables: dict,
+        excluded_types: Set[Tuple[int, int]] = set(),
+        suffix: str = "",
+    ) -> dict:
+        """Serialize network.
+
+        Parameters
+        ----------
+        ntypes : int
+            The number of types
+        ndim : int
+            The dimension of elements
+        in_dim : int
+            The input dimension
+        neuron : List[int]
+            The neuron list
+        activation_function : str
+            The activation function
+        resnet_dt : bool
+            Whether to use resnet
+        variables : dict
+            The input variables
+        excluded_types : Set[Tuple[int, int]], optional
+            The excluded types
+        suffix : str, optional
+            The suffix of the scope
+
+        Returns
+        -------
+        dict
+            The converted network data
+        """
+        embeddings = NetworkCollection(
+            ntypes=ntypes,
+            ndim=ndim,
+            network_type="embedding_network",
+        )
+        if ndim == 2:
+            for type_i, type_j in excluded_types:
+                # initialize an empty network for the excluded types
+                embeddings[(type_i, type_j)] = EmbeddingNet(
+                    in_dim=in_dim,
+                    neuron=neuron,
+                    activation_function=activation_function,
+                    resnet_dt=resnet_dt,
+                    precision=self.precision.name,
+                )
+                embeddings[(type_j, type_i)] = EmbeddingNet(
+                    in_dim=in_dim,
+                    neuron=neuron,
+                    activation_function=activation_function,
+                    resnet_dt=resnet_dt,
+                    precision=self.precision.name,
+                )
+                for layer in range(len(neuron)):
+                    embeddings[(type_i, type_j)][layer]["w"][:] = 0.0
+                    embeddings[(type_i, type_j)][layer]["b"][:] = 0.0
+                    if embeddings[(type_i, type_j)][layer]["idt"] is not None:
+                        embeddings[(type_i, type_j)][layer]["idt"][:] = 0.0
+                    embeddings[(type_j, type_i)][layer]["w"][:] = 0.0
+                    embeddings[(type_j, type_i)][layer]["b"][:] = 0.0
+                    if embeddings[(type_j, type_i)][layer]["idt"] is not None:
+                        embeddings[(type_j, type_i)][layer]["idt"][:] = 0.0
+
+        if suffix != "":
+            embedding_net_pattern = (
+                EMBEDDING_NET_PATTERN.replace("/(idt)", suffix + "/(idt)")
+                .replace("/(bias)", suffix + "/(bias)")
+                .replace("/(matrix)", suffix + "/(matrix)")
+            )
+        else:
+            embedding_net_pattern = EMBEDDING_NET_PATTERN
+        for key, value in variables.items():
+            m = re.search(embedding_net_pattern, key)
+            m = [mm for mm in m.groups() if mm is not None]
+            typei = m[0]
+            typej = "_".join(m[3:]) if len(m[3:]) else "all"
+            layer_idx = int(m[2]) - 1
+            weight_name = m[1]
+            if ndim == 0:
+                network_idx = ()
+            elif ndim == 1:
+                network_idx = (int(typej),)
+            elif ndim == 2:
+                network_idx = (int(typei), int(typej))
+            else:
+                raise ValueError(f"Invalid ndim: {ndim}")
+            if embeddings[network_idx] is None:
+                # initialize the network if it is not initialized
+                embeddings[network_idx] = EmbeddingNet(
+                    in_dim=in_dim,
+                    neuron=neuron,
+                    activation_function=activation_function,
+                    resnet_dt=resnet_dt,
+                    precision=self.precision.name,
+                )
+            assert embeddings[network_idx] is not None
+            if weight_name == "idt":
+                value = value.ravel()
+            embeddings[network_idx][layer_idx][weight_name] = value
+        return embeddings.serialize()
+
+    @classmethod
+    def deserialize_network(cls, data: dict, suffix: str = "") -> dict:
+        """Deserialize network.
+
+        Parameters
+        ----------
+        data : dict
+            The input network data
+        suffix : str, optional
+            The suffix of the scope
+
+        Returns
+        -------
+        variables : dict
+            The input variables
+        """
+        embedding_net_variables = {}
+        embeddings = NetworkCollection.deserialize(data)
+        for ii in range(embeddings.ntypes**embeddings.ndim):
+            net_idx = []
+            rest_ii = ii
+            for _ in range(embeddings.ndim):
+                net_idx.append(rest_ii % embeddings.ntypes)
+                rest_ii //= embeddings.ntypes
+            net_idx = tuple(net_idx)
+            if embeddings.ndim in (0, 1):
+                key0 = "all"
+                key1 = f"_{ii}"
+            elif embeddings.ndim == 2:
+                key0 = f"{net_idx[0]}"
+                key1 = f"_{net_idx[1]}"
+            else:
+                raise ValueError(f"Invalid ndim: {embeddings.ndim}")
+            network = embeddings[net_idx]
+            assert network is not None
+            for layer_idx, layer in enumerate(network.layers):
+                embedding_net_variables[
+                    f"filter_type_{key0}{suffix}/matrix_{layer_idx + 1}{key1}"
+                ] = layer.w
+                embedding_net_variables[
+                    f"filter_type_{key0}{suffix}/bias_{layer_idx + 1}{key1}"
+                ] = layer.b
+                if layer.idt is not None:
+                    embedding_net_variables[
+                        f"filter_type_{key0}{suffix}/idt_{layer_idx + 1}{key1}"
+                    ] = layer.idt.reshape(1, -1)
+                else:
+                    # prevent keyError
+                    embedding_net_variables[
+                        f"filter_type_{key0}{suffix}/idt_{layer_idx + 1}{key1}"
+                    ] = 0.0
+        return embedding_net_variables