refact(dpmodel,pt_expt): fitting net

deepmd/dpmodel/utils/network.py

@@ -1009,7 +1009,118 @@ def deserialize(cls, data: dict) -> "FittingNet":
     return FN
 
 
-FittingNet = make_fitting_network(EmbeddingNet, NativeNet, NativeLayer)
+class FittingNet(EmbeddingNet):
+    """The fitting network. It may be implemented as an embedding
+    net connected with a linear output layer.
+
+    Parameters
+    ----------
+    in_dim
+        Input dimension.
+    out_dim
+        Output dimension
+    neuron
+        The number of neurons in each hidden layer.
+    activation_function
+        The activation function.
+    resnet_dt
+        Use time step at the resnet architecture.
+    precision
+        Floating point precision for the model parameters.
+    bias_out
+        The last linear layer has bias.
+    seed : int, optional
+        Random seed.
+    trainable : bool or list[bool], optional
+        Whether the network is trainable.
+    """
+
+    def __init__(
+        self,
+        in_dim: int,
+        out_dim: int,
+        neuron: list[int] = [24, 48, 96],
+        activation_function: str = "tanh",
+        resnet_dt: bool = False,
+        precision: str = DEFAULT_PRECISION,
+        bias_out: bool = True,
+        seed: int | list[int] | None = None,
+        trainable: bool | list[bool] = True,
+    ) -> None:
+        if trainable is None:
+            trainable = [True] * (len(neuron) + 1)
+        elif isinstance(trainable, bool):
+            trainable = [trainable] * (len(neuron) + 1)
+        else:
+            pass
+        super().__init__(
+            in_dim,
+            neuron=neuron,
+            activation_function=activation_function,
+            resnet_dt=resnet_dt,
+            precision=precision,
+            seed=seed,
+            trainable=trainable[:-1],
+        )
+        i_in = neuron[-1] if len(neuron) > 0 else in_dim
+        i_ot = out_dim
+        self.layers.append(
+            NativeLayer(
+                i_in,
+                i_ot,
+                bias=bias_out,
+                use_timestep=False,
+                activation_function=None,
+                resnet=False,
+                precision=precision,
+                seed=child_seed(seed, len(neuron)),
+                trainable=trainable[-1],
+            )
+        )
+        self.out_dim = out_dim
+        self.bias_out = bias_out
+
+    def serialize(self) -> dict:
+        """Serialize the network to a dict.
+
+        Returns
+        -------
+        dict
+            The serialized network.
+        """
+        return {
+            "@class": "FittingNetwork",
+            "@version": 1,
+            "in_dim": self.in_dim,
+            "out_dim": self.out_dim,
+            "neuron": self.neuron.copy(),
+            "activation_function": self.activation_function,
+            "resnet_dt": self.resnet_dt,
+            "precision": self.precision,
+            "bias_out": self.bias_out,
+            "layers": [layer.serialize() for layer in self.layers],
+        }
+
+    @classmethod
+    def deserialize(cls, data: dict) -> "FittingNet":
+        """Deserialize the network from a dict.
+
+        Parameters
+        ----------
+        data : dict
+            The dict to deserialize from.
+        """
+        data = data.copy()
+        check_version_compatibility(data.pop("@version", 1), 1, 1)
+        data.pop("@class", None)
+        layers = data.pop("layers")
+        obj = cls(**data)
+        # Use type(obj.layers[0]) to respect subclass layer types
+        layer_type = type(obj.layers[0])
+        obj.layers = type(obj.layers)(
+            [layer_type.deserialize(layer) for layer in layers]
+        )
+        return obj
 
 
 class NetworkCollection:

deepmd/pt_expt/utils/network.py

@@ -11,11 +11,11 @@
     NativeOP,
 )
 from deepmd.dpmodel.utils.network import EmbeddingNet as EmbeddingNetDP
+from deepmd.dpmodel.utils.network import FittingNet as FittingNetDP
 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_fitting_network,
     make_multilayer_network,
 )
 from deepmd.pt_expt.common import (
@@ -114,8 +114,26 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
 )
 
 
-class FittingNet(make_fitting_network(EmbeddingNet, NativeNet, NativeLayer)):
-    pass
+class FittingNet(FittingNetDP, torch.nn.Module):
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
+        torch.nn.Module.__init__(self)
+        FittingNetDP.__init__(self, *args, **kwargs)
+        # Convert dpmodel layers to pt_expt NativeLayer
+        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(
+    FittingNetDP,
+    lambda v: FittingNet.deserialize(v.serialize()),
+)
 
 
 class NetworkCollection(NetworkCollectionDP, torch.nn.Module):

source/tests/common/dpmodel/test_network.py

@@ -313,6 +313,104 @@ def test_fitting_net(self) -> None:
             en1.call(inp)
             np.testing.assert_allclose(en0.call(inp), en1.call(inp))
 
+    def test_is_concrete_class(self) -> None:
+        """Verify FittingNet is a concrete class, not factory-generated."""
+        in_dim = 4
+        out_dim = 1
+        neuron = [8, 16]
+        net = FittingNet(
+            in_dim=in_dim,
+            out_dim=out_dim,
+            neuron=neuron,
+            activation_function="tanh",
+            resnet_dt=True,
+            precision="float64",
+            bias_out=True,
+        )
+        # Check it's the actual FittingNet class, not a dynamic class
+        self.assertEqual(net.__class__.__name__, "FittingNet")
+        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.out_dim, out_dim)
+        self.assertEqual(net.neuron, neuron)
+        self.assertEqual(net.activation_function, "tanh")
+        self.assertEqual(net.resnet_dt, True)
+        self.assertEqual(net.bias_out, True)
+        # FittingNet has len(neuron) embedding layers + 1 output layer
+        self.assertEqual(len(net.layers), len(neuron) + 1)
+
+    def test_forward_pass(self) -> None:
+        """Test FittingNet forward pass produces correct output shape."""
+        in_dim = 4
+        out_dim = 3
+        neuron = [8, 16, 32]
+        net = FittingNet(
+            in_dim=in_dim,
+            out_dim=out_dim,
+            neuron=neuron,
+            activation_function="tanh",
+            resnet_dt=True,
+            precision="float64",
+        )
+        # Single sample
+        rng = np.random.default_rng()
+        x = rng.standard_normal(in_dim)
+        out = net.call(x)
+        self.assertEqual(out.shape, (out_dim,))
+
+        # Batch of samples
+        batch_size = 5
+        x_batch = rng.standard_normal((batch_size, in_dim))
+        out_batch = net.call(x_batch)
+        self.assertEqual(out_batch.shape, (batch_size, out_dim))
+
+    def test_trainable_parameter_variants(self) -> None:
+        """Test FittingNet with different trainable configurations."""
+        in_dim = 4
+        out_dim = 2
+        neuron = [8, 16]
+
+        # Test 1: All layers trainable (default)
+        net_all_trainable = FittingNet(
+            in_dim=in_dim,
+            out_dim=out_dim,
+            neuron=neuron,
+            trainable=True,
+        )
+        for layer in net_all_trainable.layers:
+            self.assertTrue(layer.trainable)
+
+        # Test 2: All layers frozen
+        net_all_frozen = FittingNet(
+            in_dim=in_dim,
+            out_dim=out_dim,
+            neuron=neuron,
+            trainable=False,
+        )
+        for layer in net_all_frozen.layers:
+            self.assertFalse(layer.trainable)
+
+        # Test 3: Mixed trainable (embedding layers frozen, output layer trainable)
+        trainable_list = [False, False, True]  # 2 embedding layers + 1 output layer
+        net_mixed = FittingNet(
+            in_dim=in_dim,
+            out_dim=out_dim,
+            neuron=neuron,
+            trainable=trainable_list,
+        )
+        self.assertFalse(net_mixed.layers[0].trainable)  # First embedding layer
+        self.assertFalse(net_mixed.layers[1].trainable)  # Second embedding layer
+        self.assertTrue(net_mixed.layers[2].trainable)  # Output layer
+
+        # Test 4: Serialize/deserialize preserves trainable
+        serialized = net_mixed.serialize()
+        net_restored = FittingNet.deserialize(serialized)
+        for orig_layer, restored_layer in zip(
+            net_mixed.layers, net_restored.layers, strict=True
+        ):
+            self.assertEqual(orig_layer.trainable, restored_layer.trainable)
+
 
 class TestNetworkCollection(unittest.TestCase):
     def setUp(self) -> None:

source/tests/pt_expt/utils/test_network.py

@@ -281,3 +281,124 @@ def test_trainable_parameter_handling(self) -> None:
         for layer in net_frozen.layers:
             if layer.w is not None:
                 self.assertFalse(layer.w.requires_grad)
+
+
+class TestFittingNetRefactor(unittest.TestCase):
+    """Tests for the refactored FittingNet pt_expt wrapper."""
+
+    def setUp(self) -> None:
+        self.in_dim = 4
+        self.out_dim = 1
+        self.neuron = [8, 16]
+        self.activation = "tanh"
+        self.resnet_dt = True
+        self.precision = "float64"
+
+    def test_pt_expt_fitting_net_wraps_dpmodel(self) -> None:
+        """Verify pt_expt FittingNet correctly wraps dpmodel."""
+        from deepmd.pt_expt.utils.network import (
+            FittingNet,
+        )
+
+        net = FittingNet(
+            in_dim=self.in_dim,
+            out_dim=self.out_dim,
+            neuron=self.neuron,
+            activation_function=self.activation,
+            resnet_dt=self.resnet_dt,
+            precision=self.precision,
+            seed=GLOBAL_SEED,
+        )
+        # Check it's a torch.nn.Module
+        self.assertIsInstance(net, torch.nn.Module)
+        # Check layers are converted to pt_expt NativeLayer (torch modules)
+        self.assertIsInstance(net.layers, torch.nn.ModuleList)
+        for layer in net.layers:
+            self.assertIsInstance(layer, torch.nn.Module)
+
+    def test_pt_expt_fitting_net_forward(self) -> None:
+        """Test pt_expt FittingNet forward pass returns torch.Tensor."""
+        from deepmd.pt_expt.utils.network import (
+            FittingNet,
+        )
+
+        net = FittingNet(
+            in_dim=self.in_dim,
+            out_dim=self.out_dim,
+            neuron=self.neuron,
+            activation_function=self.activation,
+            resnet_dt=self.resnet_dt,
+            precision=self.precision,
+            seed=GLOBAL_SEED,
+        )
+        x = torch.randn(5, self.in_dim, dtype=torch.float64, device=env.DEVICE)
+        out = net(x)
+        self.assertIsInstance(out, torch.Tensor)
+        self.assertEqual(out.shape, (5, self.out_dim))
+        self.assertEqual(out.dtype, torch.float64)
+
+    def test_serialization_round_trip_pt_expt(self) -> None:
+        """Test pt_expt FittingNet serialization/deserialization."""
+        from deepmd.pt_expt.utils.network import (
+            FittingNet,
+        )
+
+        net = FittingNet(
+            in_dim=self.in_dim,
+            out_dim=self.out_dim,
+            neuron=self.neuron,
+            activation_function=self.activation,
+            resnet_dt=self.resnet_dt,
+            precision=self.precision,
+            seed=GLOBAL_SEED,
+        )
+        x = torch.randn(5, self.in_dim, dtype=torch.float64, device=env.DEVICE)
+        out1 = net(x)
+
+        # Serialize and deserialize
+        serialized = net.serialize()
+        net2 = FittingNet.deserialize(serialized)
+
+        # Verify layers are still pt_expt NativeLayer modules
+        self.assertIsInstance(net2.layers, torch.nn.ModuleList)
+        for layer in net2.layers:
+            self.assertIsInstance(layer, torch.nn.Module)
+
+        out2 = net2(x)
+        np.testing.assert_allclose(
+            out1.detach().cpu().numpy(),
+            out2.detach().cpu().numpy(),
+        )
+
+    def test_registry_converts_dpmodel_to_pt_expt(self) -> None:
+        """Test that dpmodel FittingNet can be converted to pt_expt via registry."""
+        from deepmd.dpmodel.utils.network import FittingNet as DPFittingNet
+        from deepmd.pt_expt.common import (
+            try_convert_module,
+        )
+        from deepmd.pt_expt.utils.network import (
+            FittingNet,
+        )
+
+        # Create dpmodel FittingNet
+        dp_net = DPFittingNet(
+            in_dim=self.in_dim,
+            out_dim=self.out_dim,
+            neuron=self.neuron,
+            activation_function=self.activation,
+            resnet_dt=self.resnet_dt,
+            precision=self.precision,
+            seed=GLOBAL_SEED,
+        )
+
+        # Try to convert via registry
+        converted = try_convert_module(dp_net)
+
+        # Should return pt_expt FittingNet
+        self.assertIsNotNone(converted)
+        self.assertIsInstance(converted, torch.nn.Module)
+        self.assertIsInstance(converted, FittingNet)
+
+        # Verify layers are pt_expt modules
+        for layer in converted.layers:
+            self.assertIsInstance(layer, torch.nn.Module)