fix se_atten compression when suffix is given

Signed-off-by: Jinzhe Zeng <[email protected]>
njzjz · Jun 25, 2023 · 105086c · 105086c
1 parent 089b1b1
commit 105086c
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 6 deletions.
diff --git a/deepmd/descriptor/se_atten.py b/deepmd/descriptor/se_atten.py
@@ -429,11 +429,9 @@ def _get_two_side_embedding_net_variable(self, graph_def, varialbe_name, suffix)
         for i in range(1, self.layer_size + 1):
             target = get_pattern_nodes_from_graph_def(
                 graph_def,
-                f"filter_type_all{suffix}/{varialbe_name}_{i}{suffix}_two_side_ebd",
+                f"filter_type_all{suffix}/{varialbe_name}_{i}_two_side_ebd",
             )
-            node = target[
-                f"filter_type_all{suffix}/{varialbe_name}_{i}{suffix}_two_side_ebd"
-            ]
+            node = target[f"filter_type_all{suffix}/{varialbe_name}_{i}_two_side_ebd"]
             ret["layer_" + str(i)] = node
         return ret
 

diff --git a/deepmd/utils/tabulate.py b/deepmd/utils/tabulate.py
@@ -420,7 +420,7 @@ def _get_bias(self):
             bias["layer_" + str(layer)] = []
             if isinstance(self.descrpt, deepmd.descriptor.DescrptSeAtten):
                 node = self.embedding_net_nodes[
-                    f"filter_type_all{self.suffix}/bias_{layer}{self.suffix}"
+                    f"filter_type_all{self.suffix}/bias_{layer}"
                 ]
                 bias["layer_" + str(layer)].append(tf.make_ndarray(node))
             elif isinstance(self.descrpt, deepmd.descriptor.DescrptSeA):
@@ -484,7 +484,7 @@ def _get_matrix(self):
             matrix["layer_" + str(layer)] = []
             if isinstance(self.descrpt, deepmd.descriptor.DescrptSeAtten):
                 node = self.embedding_net_nodes[
-                    f"filter_type_all{self.suffix}/matrix_{layer}{self.suffix}"
+                    f"filter_type_all{self.suffix}/matrix_{layer}"
                 ]
                 matrix["layer_" + str(layer)].append(tf.make_ndarray(node))
             elif isinstance(self.descrpt, deepmd.descriptor.DescrptSeA):