fix(tf): fix compress suffix in DescrptDPA1Compat (#4243)

Fix #4114 .
<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

- **New Features**
- Enhanced compression capabilities in descriptor models with new
optional parameters for improved flexibility.
- Improved serialization processes for attention layers, allowing for
better handling of scaling factors and normalization.
- Dynamic tensor name construction in utility functions to accommodate
varying suffixes.

- **Bug Fixes**
- Adjusted method parameters to ensure compatibility and functionality
with new suffix options.

- **Tests**
- Introduced a new test suite to validate the functionality of the
TensorFlow-based descriptor model, ensuring consistent output with the
updated features.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: iProzd

deepmd/tf/descriptor/se_atten.py

@@ -423,6 +423,7 @@ def enable_compression(
         table_stride_2: float = 0.1,
         check_frequency: int = -1,
         suffix: str = "",
+        tebd_suffix: str = "",
     ) -> None:
         """Reveive the statisitcs (distance, max_nbor_size and env_mat_range) of the training data.
 
@@ -444,6 +445,8 @@ def enable_compression(
             The overflow check frequency
         suffix : str, optional
             The suffix of the scope
+        tebd_suffix : str, optional
+            The suffix of the type embedding scope, only for DescrptDPA1Compat
         """
         # do some checks before the mocel compression process
         assert (
@@ -496,7 +499,9 @@ def enable_compression(
             min_nbor_dist, table_extrapolate, table_stride_1, table_stride_2
         )
 
-        self.final_type_embedding = get_two_side_type_embedding(self, graph)
+        self.final_type_embedding = get_two_side_type_embedding(
+            self, graph, suffix=tebd_suffix
+        )
         type_side_suffix = get_extra_embedding_net_suffix(type_one_side=False)
         self.matrix = get_extra_side_embedding_net_variable(
             self, graph_def, type_side_suffix, "matrix", suffix
@@ -2248,6 +2253,56 @@ def build(
             self.dout = tf.concat([self.dout, atom_embed], axis=-1)
         return self.dout
 
+    def enable_compression(
+        self,
+        min_nbor_dist: float,
+        graph: tf.Graph,
+        graph_def: tf.GraphDef,
+        table_extrapolate: float = 5,
+        table_stride_1: float = 0.01,
+        table_stride_2: float = 0.1,
+        check_frequency: int = -1,
+        suffix: str = "",
+        tebd_suffix: str = "",
+    ) -> None:
+        """Reveive the statisitcs (distance, max_nbor_size and env_mat_range) of the training data.
+
+        Parameters
+        ----------
+        min_nbor_dist
+            The nearest distance between atoms
+        graph : tf.Graph
+            The graph of the model
+        graph_def : tf.GraphDef
+            The graph_def of the model
+        table_extrapolate
+            The scale of model extrapolation
+        table_stride_1
+            The uniform stride of the first table
+        table_stride_2
+            The uniform stride of the second table
+        check_frequency
+            The overflow check frequency
+        suffix : str, optional
+            The suffix of the scope
+        tebd_suffix : str, optional
+            Same as suffix.
+        """
+        assert (
+            tebd_suffix == ""
+        ), "DescrptDPA1Compat must use the same tebd_suffix as suffix!"
+        super().enable_compression(
+            min_nbor_dist,
+            graph,
+            graph_def,
+            table_extrapolate=table_extrapolate,
+            table_stride_1=table_stride_1,
+            table_stride_2=table_stride_2,
+            check_frequency=check_frequency,
+            suffix=suffix,
+            tebd_suffix=suffix,
+        )
+
     def init_variables(
         self,
         graph: tf.Graph,

deepmd/tf/utils/compress.py

@@ -20,8 +20,8 @@ def get_type_embedding(self, graph):
     return type_embedding
 
 
-def get_two_side_type_embedding(self, graph):
-    type_embedding = get_tensor_by_name_from_graph(graph, "t_typeebd")
+def get_two_side_type_embedding(self, graph, suffix=""):
+    type_embedding = get_tensor_by_name_from_graph(graph, f"t_typeebd{suffix}")
     type_embedding = type_embedding.astype(self.filter_np_precision)
     type_embedding_shape = type_embedding.shape
 

deepmd/tf/utils/tabulate.py

@@ -126,7 +126,7 @@ def __init__(
         self.dstd = get_tensor_by_name_from_graph(
             self.graph, f"descrpt_attr{self.suffix}/t_std"
         )
-        self.ntypes = get_tensor_by_name_from_graph(self.graph, "descrpt_attr/ntypes")
+        self.ntypes = self.descrpt.get_ntypes()
 
         self.embedding_net_nodes = get_embedding_net_nodes_from_graph_def(
             self.graph_def, suffix=self.suffix

source/tests/tf/test_model_compression_dpa1_compat_suffix_only.py

@@ -0,0 +1,153 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+import unittest
+
+import numpy as np
+
+from deepmd.common import (
+    make_default_mesh,
+)
+from deepmd.env import (
+    GLOBAL_NP_FLOAT_PRECISION,
+)
+from deepmd.tf.descriptor.se_atten import DescrptDPA1Compat as tf_SeAtten
+from deepmd.tf.env import (
+    GLOBAL_TF_FLOAT_PRECISION,
+    default_tf_session_config,
+    tf,
+)
+from deepmd.tf.utils.sess import (
+    run_sess,
+)
+
+
+def build_tf_descriptor(obj, natoms, coords, atype, box, suffix):
+    t_coord = tf.placeholder(GLOBAL_TF_FLOAT_PRECISION, [None], name="i_coord")
+    t_type = tf.placeholder(tf.int32, [None], name="i_type")
+    t_natoms = tf.placeholder(tf.int32, natoms.shape, name="i_natoms")
+    t_box = tf.placeholder(GLOBAL_TF_FLOAT_PRECISION, [9], name="i_box")
+    t_mesh = tf.placeholder(tf.int32, [None], name="i_mesh")
+    t_des = obj.build(
+        t_coord,
+        t_type,
+        t_natoms,
+        t_box,
+        t_mesh,
+        {},
+        suffix=suffix,
+    )
+    return [t_des], {
+        t_coord: coords,
+        t_type: atype,
+        t_natoms: natoms,
+        t_box: box,
+        t_mesh: make_default_mesh(True, False),
+    }
+
+
+def build_eval_tf(sess, obj, natoms, coords, atype, box, suffix):
+    t_out, feed_dict = build_tf_descriptor(obj, natoms, coords, atype, box, suffix)
+
+    t_out_indentity = [
+        tf.identity(tt, name=f"o_{ii}_{suffix}") for ii, tt in enumerate(t_out)
+    ]
+    run_sess(sess, tf.global_variables_initializer())
+    return run_sess(
+        sess,
+        t_out_indentity,
+        feed_dict=feed_dict,
+    )
+
+
+class TestDescriptorSeA(unittest.TestCase):
+    def setUp(self):
+        self.device = "cpu"
+        self.seed = 21
+        self.sel = [9, 10]
+        self.rcut_smth = 5.80
+        self.rcut = 6.00
+        self.neuron = [6, 12, 24]
+        self.axis_neuron = 3
+        self.ntypes = 2
+        self.coords = np.array(
+            [
+                12.83,
+                2.56,
+                2.18,
+                12.09,
+                2.87,
+                2.74,
+                00.25,
+                3.32,
+                1.68,
+                3.36,
+                3.00,
+                1.81,
+                3.51,
+                2.51,
+                2.60,
+                4.27,
+                3.22,
+                1.56,
+            ],
+            dtype=GLOBAL_NP_FLOAT_PRECISION,
+        )
+        self.atype = np.array([0, 1, 1, 0, 1, 1], dtype=np.int32)
+        # self.atype = np.array([0, 0, 1, 1, 1, 1], dtype=np.int32)
+        self.box = np.array(
+            [13.0, 0.0, 0.0, 0.0, 13.0, 0.0, 0.0, 0.0, 13.0],
+            dtype=GLOBAL_NP_FLOAT_PRECISION,
+        )
+        self.natoms = np.array([6, 6, 2, 4], dtype=np.int32)
+        self.suffix = "test"
+        self.type_one_side = False
+        self.se_a_tf = tf_SeAtten(
+            self.rcut,
+            self.rcut_smth,
+            self.sel,
+            self.ntypes,
+            self.neuron,
+            self.axis_neuron,
+            type_one_side=self.type_one_side,
+            seed=21,
+            precision="float32",
+            tebd_input_mode="strip",
+            temperature=1.0,
+            attn_layer=0,
+        )
+
+    def test_tf_pt_consistent(
+        self,
+    ):
+        with tf.Session(config=default_tf_session_config) as sess:
+            graph = tf.get_default_graph()
+            ret = build_eval_tf(
+                sess,
+                self.se_a_tf,
+                self.natoms,
+                self.coords,
+                self.atype,
+                self.box,
+                self.suffix,
+            )
+            output_graph_def = tf.graph_util.convert_variables_to_constants(
+                sess,
+                graph.as_graph_def(),
+                [f"o_{ii}_{self.suffix}" for ii, _ in enumerate(ret)],
+            )
+            with tf.Graph().as_default() as new_graph:
+                tf.import_graph_def(output_graph_def, name="")
+            self.se_a_tf.init_variables(
+                new_graph,
+                output_graph_def,
+                suffix=self.suffix,
+            )
+            self.se_a_tf.enable_compression(
+                1.0,
+                new_graph,
+                output_graph_def,
+                suffix=self.suffix,
+            )
+
+
+if __name__ == "__main__":
+    unittest.main()