model format for the embedding net

deepmd_utils/model_format/__init__.py

@@ -0,0 +1,9 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+from .network import (
+    EmbeddingNet,
+    NativeLayer,
+    NativeNet,
+    load_dp_model,
+    save_dp_model,
+    traverse_model_dict,
+)

deepmd_utils/model_format.py → deepmd_utils/model_format/network.py

@@ -275,10 +275,10 @@ class NativeNet(NativeOP):
         The layers of the network.
     """
 
-    def __init__(self, layers: Optional[List[NativeLayer]] = None) -> None:
+    def __init__(self, layers: Optional[List[dict]] = None) -> None:
         if layers is None:
             layers = []
-        self.layers = layers
+        self.layers = [NativeLayer.deserialize(layer) for layer in layers]
 
     def serialize(self) -> dict:
         """Serialize the network to a dict.
@@ -299,7 +299,7 @@ def deserialize(cls, data: dict) -> "NativeNet":
         data : dict
             The dict to deserialize from.
         """
-        return cls([NativeLayer.deserialize(layer) for layer in data["layers"]])
+        return cls(data["layers"])
 
     def __getitem__(self, key):
         assert isinstance(key, int)
@@ -329,3 +329,62 @@ def call(self, x: np.ndarray) -> np.ndarray:
         for layer in self.layers:
             x = layer.call(x)
         return x
+
+
+class EmbeddingNet(NativeNet):
+    def __init__(
+        self,
+        in_dim,
+        neuron: List[int] = [24, 48, 96],
+        activation_function: str = "tanh",
+        resnet_dt: bool = False,
+    ):
+        layers = []
+        i_in = in_dim
+        for idx, ii in enumerate(neuron):
+            i_ot = ii
+            layers.append(
+                NativeLayer(
+                    np.random.normal(size=(i_in, i_ot)),
+                    b=np.random.normal(size=(ii)),
+                    idt=np.random.normal(size=(ii)) if resnet_dt else None,
+                    activation_function=activation_function,
+                    resnet=True,
+                ).serialize()
+            )
+            i_in = i_ot
+        super(EmbeddingNet, self).__init__(layers)
+        self.in_dim = in_dim
+        self.neuron = neuron
+        self.activation_function = activation_function
+        self.resnet_dt = resnet_dt
+
+    def serialize(self) -> dict:
+        """Serialize the network to a dict.
+
+        Returns
+        -------
+        dict
+            The serialized network.
+        """
+        return {
+            "in_dim": self.in_dim,
+            "neuron": self.neuron,
+            "activation_function": self.activation_function,
+            "resnet_dt": self.resnet_dt,
+            "layers": [layer.serialize() for layer in self.layers],
+        }
+
+    @classmethod
+    def deserialize(cls, data: dict) -> "EmbeddingNet":
+        """Deserialize the network from a dict.
+
+        Parameters
+        ----------
+        data : dict
+            The dict to deserialize from.
+        """
+        layers = data.pop("layers")
+        obj = cls(**data)
+        super(EmbeddingNet, obj).__init__(layers)
+        return obj

source/tests/test_model_format_utils.py

@@ -9,6 +9,7 @@
 import numpy as np
 
 from deepmd_utils.model_format import (
+    EmbeddingNet,
     NativeLayer,
     NativeNet,
     load_dp_model,
@@ -90,6 +91,17 @@ def test_deserialize(self):
         np.testing.assert_array_equal(network[0]["resnet"], True)
         np.testing.assert_array_equal(network[1]["resnet"], True)
 
+    def test_embedding_net(self):
+        for ni, idt, act in itertools.product(
+            [1, 10],
+            [True, False],
+            ["tanh", "none"],
+        ):
+            en0 = EmbeddingNet(ni)
+            en1 = EmbeddingNet.deserialize(en0.serialize())
+            inp = np.ones([ni])
+            np.testing.assert_allclose(en0.call(inp), en1.call(inp))
+
 
 class TestDPModel(unittest.TestCase):
     def setUp(self) -> None: