add ut

Author: iProzd

deepmd/pt/infer/deep_eval.py

@@ -130,7 +130,6 @@ def __init__(
                         ] = state_dict[item].clone()
                 state_dict = state_dict_head
             model = get_model(self.input_param).to(DEVICE)
-            # TODO fix jit
             if not self.input_param.get("hessian_mode"):
                 model = torch.jit.script(model)
             self.dp = ModelWrapper(model)

deepmd/pt/model/task/fitting.py

@@ -617,9 +617,22 @@ def _forward_common(
             )  # Shape is [nframes, natoms[0], net_dim_out]
         else:
             if self.eval_return_middle_output:
-                raise NotImplementedError(
-                    "Middle output is only supported for mixed types!"
-                )
+                outs_middle = torch.zeros(
+                    (nf, nloc, self.neuron[-1]),
+                    dtype=self.prec,
+                    device=descriptor.device,
+                )  # jit assertion
+                for type_i, ll in enumerate(self.filter_layers.networks):
+                    mask = (atype == type_i).unsqueeze(-1)
+                    mask = torch.tile(mask, (1, 1, net_dim_out))
+                    middle_output_type = ll.call_until_last(xx)
+                    middle_output_type = torch.where(
+                        torch.tile(mask, (1, 1, self.neuron[-1])),
+                        middle_output_type,
+                        0.0,
+                    )
+                    outs_middle = outs_middle + middle_output_type
+                results["middle_output"] = outs_middle
             for type_i, ll in enumerate(self.filter_layers.networks):
                 mask = (atype == type_i).unsqueeze(-1)
                 mask = torch.tile(mask, (1, 1, net_dim_out))

source/tests/infer/case.py

@@ -125,6 +125,13 @@ def __init__(self, data: dict) -> None:
         else:
             self.descriptor = None
 
+        if "fit_ll" in data:
+            self.fit_ll = np.array(data["fit_ll"], dtype=np.float64).reshape(
+                self.nloc, -1
+            )
+        else:
+            self.fit_ll = None
+
 
 class Case:
     """Test case.

source/tests/infer/deeppot-testcase.yaml

@@ -350,6 +350,45 @@ results:
         1.391094495316195001e+00,
         7.036614101584164338e-01,
       ]
+    fit_ll:
+      [
+        -1.930622006643730598e-02,
+        7.105172146387829235e-01,
+        8.063835335367619539e-01,
+        -8.414936892447275607e-01,
+        1.076881365346436414e+00,
+        -5.058153291569045251e-01,
+        -3.104797373867691779e-02,
+        7.915138025598530414e-01,
+        8.704498369678651537e-01,
+        -9.394329433114724237e-01,
+        1.081177674358831053e+00,
+        -5.122829163516022799e-01,
+        5.307913125575804136e-03,
+        7.644783775007328863e-01,
+        8.548853566716824171e-01,
+        -9.264496186379944653e-01,
+        1.087178488222722672e+00,
+        -4.893627623467682874e-01,
+        -1.098746804357388085e-01,
+        8.092546382430507723e-01,
+        8.757043853926992361e-01,
+        -9.036627000544070754e-01,
+        1.064706190677472852e+00,
+        -5.670533963064982030e-01,
+        -1.270062329805081158e-01,
+        8.618261193779762630e-01,
+        8.979592934126284787e-01,
+        -9.939941754957831721e-01,
+        1.072078883192923771e+00,
+        -5.780043831847785363e-01,
+        -8.617331266742107865e-02,
+        8.388158674801169390e-01,
+        8.904977456468012864e-01,
+        -9.751383339999978306e-01,
+        1.075378146084344344e+00,
+        -5.508880199511664300e-01,
+      ]
   - coord:
       [
         12.83,

source/tests/infer/test_models.py

@@ -164,6 +164,24 @@ def test_descriptor(self) -> None:
             expected_descpt = result.descriptor
             np.testing.assert_almost_equal(descpt.ravel(), expected_descpt.ravel())
 
+    def test_fitting_last_layer(self) -> None:
+        _, extension = self.param
+        if extension == ".pb":
+            self.skipTest("fitting_last_layer not supported for TensorFlow models")
+        for ii, result in enumerate(self.case.results):
+            if result.fit_ll is None:
+                continue
+            fit_ll = self.dp.eval_fitting_last_layer(
+                result.coord, result.box, result.atype
+            )
+            expected_fit_ll = result.fit_ll
+            np.testing.assert_almost_equal(fit_ll.ravel(), expected_fit_ll.ravel())
+            fit_ll = self.dp.eval_fitting_last_layer(
+                result.coord, result.box, result.atype
+            )
+            expected_fit_ll = result.fit_ll
+            np.testing.assert_almost_equal(fit_ll.ravel(), expected_fit_ll.ravel())
+
     def test_2frame_atm(self) -> None:
         for ii, result in enumerate(self.case.results):
             coords2 = np.concatenate((result.coord, result.coord))