chore(dpmodel,pt_expt): refactorize the implementation of embedding net

Author: Han Wang

deepmd/dpmodel/utils/network.py

@@ -788,7 +788,112 @@ def deserialize(cls, data: dict) -> "EmbeddingNet":
     return EN
 
 
-EmbeddingNet = make_embedding_network(NativeNet, NativeLayer)
+class EmbeddingNet(NativeNet):
+    """The embedding network.
+
+    Parameters
+    ----------
+    in_dim
+        Input dimension.
+    neuron
+        The number of neurons in each layer. The output dimension
+        is the same as the dimension of the last layer.
+    activation_function
+        The activation function.
+    resnet_dt
+        Use time step at the resnet architecture.
+    precision
+        Floating point precision for the model parameters.
+    seed : int, optional
+        Random seed.
+    bias : bool, Optional
+        Whether to use bias in the embedding layer.
+    trainable : bool or list[bool], Optional
+        Whether the weights are trainable. If a list, each element
+        corresponds to a layer.
+    """
+
+    def __init__(
+        self,
+        in_dim: int,
+        neuron: list[int] = [24, 48, 96],
+        activation_function: str = "tanh",
+        resnet_dt: bool = False,
+        precision: str = DEFAULT_PRECISION,
+        seed: int | list[int] | None = None,
+        bias: bool = True,
+        trainable: bool | list[bool] = True,
+    ) -> None:
+        layers = []
+        i_in = in_dim
+        if isinstance(trainable, bool):
+            trainable = [trainable] * len(neuron)
+        for idx, ii in enumerate(neuron):
+            i_ot = ii
+            layers.append(
+                NativeLayer(
+                    i_in,
+                    i_ot,
+                    bias=bias,
+                    use_timestep=resnet_dt,
+                    activation_function=activation_function,
+                    resnet=True,
+                    precision=precision,
+                    seed=child_seed(seed, idx),
+                    trainable=trainable[idx],
+                ).serialize()
+            )
+            i_in = i_ot
+        super().__init__(layers)
+        self.in_dim = in_dim
+        self.neuron = neuron
+        self.activation_function = activation_function
+        self.resnet_dt = resnet_dt
+        self.precision = precision
+        self.bias = bias
+
+    def serialize(self) -> dict:
+        """Serialize the network to a dict.
+
+        Returns
+        -------
+        dict
+            The serialized network.
+        """
+        return {
+            "@class": "EmbeddingNetwork",
+            "@version": 2,
+            "in_dim": self.in_dim,
+            "neuron": self.neuron.copy(),
+            "activation_function": self.activation_function,
+            "resnet_dt": self.resnet_dt,
+            "bias": self.bias,
+            # make deterministic
+            "precision": np.dtype(PRECISION_DICT[self.precision]).name,
+            "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.
+        """
+        data = data.copy()
+        check_version_compatibility(data.pop("@version", 1), 2, 1)
+        data.pop("@class", None)
+        layers = data.pop("layers")
+        obj = cls(**data)
+        # Reinitialize layers from serialized data, using the same layer type
+        # that __init__ created (respects subclass overrides via MRO).
+        layer_type = type(obj.layers[0])
+        obj.layers = type(obj.layers)(
+            [layer_type.deserialize(layer) for layer in layers]
+        )
+        return obj
 
 
 def make_fitting_network(

deepmd/pt_expt/utils/network.py

@@ -10,11 +10,11 @@
 from deepmd.dpmodel.common import (
     NativeOP,
 )
+from deepmd.dpmodel.utils.network import EmbeddingNet as EmbeddingNetDP
 from deepmd.dpmodel.utils.network import LayerNorm as LayerNormDP
 from deepmd.dpmodel.utils.network import NativeLayer as NativeLayerDP
 from deepmd.dpmodel.utils.network import NetworkCollection as NetworkCollectionDP
 from deepmd.dpmodel.utils.network import (
-    make_embedding_network,
     make_fitting_network,
     make_multilayer_network,
 )
@@ -91,8 +91,27 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
         return self.call(x)
 
 
-class EmbeddingNet(make_embedding_network(NativeNet, NativeLayer)):
-    pass
+class EmbeddingNet(EmbeddingNetDP, torch.nn.Module):
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
+        torch.nn.Module.__init__(self)
+        EmbeddingNetDP.__init__(self, *args, **kwargs)
+        # EmbeddingNetDP.__init__ creates dpmodel NativeLayer instances.
+        # Convert to pt_expt NativeLayer and wrap in ModuleList.
+        self.layers = torch.nn.ModuleList(
+            [NativeLayer.deserialize(layer.serialize()) for layer in self.layers]
+        )
+
+    def __call__(self, *args: Any, **kwargs: Any) -> Any:
+        return torch.nn.Module.__call__(self, *args, **kwargs)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.call(x)
+
+
+register_dpmodel_mapping(
+    EmbeddingNetDP,
+    lambda v: EmbeddingNet.deserialize(v.serialize()),
+)
 
 
 class FittingNet(make_fitting_network(EmbeddingNet, NativeNet, NativeLayer)):

source/tests/common/dpmodel/test_network.py

@@ -180,6 +180,76 @@ def test_embedding_net(self) -> None:
             inp = np.ones([ni], dtype=get_xp_precision(np, prec))
             np.testing.assert_allclose(en0.call(inp), en1.call(inp))
 
+    def test_is_concrete_class(self) -> None:
+        """Verify EmbeddingNet is a concrete class, not factory-generated."""
+        in_dim = 4
+        neuron = [8, 16, 32]
+        net = EmbeddingNet(
+            in_dim=in_dim,
+            neuron=neuron,
+            activation_function="tanh",
+            resnet_dt=True,
+            precision="float64",
+        )
+        # Check it's the actual EmbeddingNet class, not a dynamic class
+        self.assertEqual(net.__class__.__name__, "EmbeddingNet")
+        self.assertEqual(net.__class__.__module__, "deepmd.dpmodel.utils.network")
+        # Verify it has the expected attributes
+        self.assertEqual(net.in_dim, in_dim)
+        self.assertEqual(net.neuron, neuron)
+        self.assertEqual(net.activation_function, "tanh")
+        self.assertEqual(net.resnet_dt, True)
+        self.assertEqual(len(net.layers), len(neuron))
+
+    def test_forward_pass(self) -> None:
+        """Test EmbeddingNet forward pass produces correct shapes."""
+        in_dim = 4
+        neuron = [8, 16, 32]
+        net = EmbeddingNet(
+            in_dim=in_dim,
+            neuron=neuron,
+            activation_function="tanh",
+            resnet_dt=True,
+            precision="float64",
+        )
+        rng = np.random.default_rng()
+        x = rng.standard_normal((5, in_dim))
+        out = net.call(x)
+        self.assertEqual(out.shape, (5, neuron[-1]))
+        self.assertEqual(out.dtype, np.float64)
+
+    def test_trainable_parameter_variants(self) -> None:
+        """Test EmbeddingNet with different trainable configurations."""
+        in_dim = 4
+        neuron = [8, 16]
+
+        # All trainable
+        net_trainable = EmbeddingNet(
+            in_dim=in_dim,
+            neuron=neuron,
+            trainable=True,
+        )
+        for layer in net_trainable.layers:
+            self.assertTrue(layer.trainable)
+
+        # All frozen
+        net_frozen = EmbeddingNet(
+            in_dim=in_dim,
+            neuron=neuron,
+            trainable=False,
+        )
+        for layer in net_frozen.layers:
+            self.assertFalse(layer.trainable)
+
+        # Mixed trainable
+        net_mixed = EmbeddingNet(
+            in_dim=in_dim,
+            neuron=neuron,
+            trainable=[True, False],
+        )
+        self.assertTrue(net_mixed.layers[0].trainable)
+        self.assertFalse(net_mixed.layers[1].trainable)
+
 
 class TestFittingNet(unittest.TestCase):
     def test_fitting_net(self) -> None: