add model format for the embedding net

deepmd_utils/model_format/__init__.py

@@ -0,0 +1,8 @@
+from .network import (
+  traverse_model_dict,
+  save_dp_model,
+  load_dp_model,
+  NativeLayer,
+  NativeNet,
+  EmbeddingNet,
+)

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,61 @@ 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

@@ -11,6 +11,7 @@
 from deepmd_utils.model_format import (
     NativeLayer,
     NativeNet,
+    EmbeddingNet,
     load_dp_model,
     save_dp_model,
 )
@@ -91,6 +92,18 @@ def test_deserialize(self):
         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:
         self.w = np.full((3, 2), 3.0)