make compress work for hybrid descriptor composed of se_e2_a (#1094)

* make compress work for hybrid descriptor composed of se_e2_a

* fix `get_embedding_net_nodes_from_graph_def`

* fix lint warning

Author: njzjz

deepmd/descriptor/descriptor.py

@@ -188,7 +188,8 @@ def enable_compression(self,
                            table_extrapolate: float = 5.,
                            table_stride_1: float = 0.01,
                            table_stride_2: float = 0.1,
-                           check_frequency: int = -1
+                           check_frequency: int = -1,
+                           suffix: str = "",
                            ) -> None:
         """
         Reveive the statisitcs (distance, max_nbor_size and env_mat_range) of the
@@ -208,13 +209,15 @@ def enable_compression(self,
                 The uniform stride of the second table
         check_frequency : int, default: -1
                 The overflow check frequency
+        suffix : str, optional
+                The suffix of the scope
 
         Notes
         -----
         This method is called by others when the descriptor supported compression.
         """
         raise NotImplementedError(
-            "Descriptor %s doesn't support compression!" % self.__name__)
+            "Descriptor %s doesn't support compression!" % type(self).__name__)
 
     @abstractmethod
     def prod_force_virial(self,

deepmd/descriptor/hybrid.py

@@ -220,3 +220,36 @@ def prod_force_virial(self,
                 virial += vv
                 atom_virial += av
         return force, virial, atom_virial
+
+    def enable_compression(self,
+                           min_nbor_dist: float,
+                           model_file: str = 'frozon_model.pb',
+                           table_extrapolate: float = 5.,
+                           table_stride_1: float = 0.01,
+                           table_stride_2: float = 0.1,
+                           check_frequency: int = -1,
+                           suffix: str = ""
+                           ) -> None:
+        """
+        Reveive the statisitcs (distance, max_nbor_size and env_mat_range) of the
+        training data.
+
+        Parameters
+        ----------
+        min_nbor_dist : float
+                The nearest distance between atoms
+        model_file : str, default: 'frozon_model.pb'
+                The original frozen model, which will be compressed by the program
+        table_extrapolate : float, default: 5.
+                The scale of model extrapolation
+        table_stride_1 : float, default: 0.01
+                The uniform stride of the first table
+        table_stride_2 : float, default: 0.1
+                The uniform stride of the second table
+        check_frequency : int, default: -1
+                The overflow check frequency
+        suffix : str, optional
+                The suffix of the scope
+        """
+        for idx, ii in enumerate(self.descrpt_list):
+            ii.enable_compression(min_nbor_dist, model_file, table_extrapolate, table_stride_1, table_stride_2, check_frequency, suffix=f"{suffix}_{idx}")

deepmd/descriptor/se_a.py

@@ -297,7 +297,8 @@ def enable_compression(self,
                            table_extrapolate : float = 5,
                            table_stride_1 : float = 0.01,
                            table_stride_2 : float = 0.1,
-                           check_frequency : int = -1
+                           check_frequency : int = -1,
+                           suffix : str = "",
     ) -> None:
         """
         Reveive the statisitcs (distance, max_nbor_size and env_mat_range) of the training data.
@@ -316,10 +317,15 @@ def enable_compression(self,
                 The uniform stride of the second table
         check_frequency
                 The overflow check frequency
+        suffix : str, optional
+                The suffix of the scope
         """
+        assert (
+            not self.filter_resnet_dt
+        ), "Model compression error: descriptor resnet_dt must be false!"
         self.compress = True
         self.table = DPTabulate(
-            model_file, self.type_one_side, self.exclude_types, self.compress_activation_fn)
+            model_file, self.type_one_side, self.exclude_types, self.compress_activation_fn, suffix=suffix)
         self.table_config = [table_extrapolate, table_stride_1, table_stride_2, check_frequency]
         self.lower, self.upper \
             = self.table.build(min_nbor_dist, 
@@ -328,8 +334,8 @@ def enable_compression(self,
                                table_stride_2)
 
         graph, _ = load_graph_def(model_file)
-        self.davg = get_tensor_by_name_from_graph(graph, 'descrpt_attr/t_avg')
-        self.dstd = get_tensor_by_name_from_graph(graph, 'descrpt_attr/t_std')
+        self.davg = get_tensor_by_name_from_graph(graph, 'descrpt_attr%s/t_avg' % suffix)
+        self.dstd = get_tensor_by_name_from_graph(graph, 'descrpt_attr%s/t_std' % suffix)
 
 
 

deepmd/entrypoints/compress.py

@@ -105,12 +105,7 @@ def compress(
     jdata = normalize(jdata)
 
     # check the descriptor info of the input file
-    assert (
-        jdata["model"]["descriptor"]["type"] == "se_a" or jdata["model"]["descriptor"]["type"] == "se_e2_a"
-    ), "Model compression error: descriptor type must be se_a or se_e2_a!"
-    assert (
-        jdata["model"]["descriptor"]["resnet_dt"] is False
-    ), "Model compression error: descriptor resnet_dt must be false!"
+    # move to the specific Descriptor class
 
     # stage 1: training or refining the model with tabulation
     log.info("\n\n")

deepmd/train/trainer.py

@@ -325,7 +325,6 @@ def build (self,
                     name = 'train_attr/max_nbor_size',
                     dtype = GLOBAL_TF_FLOAT_PRECISION)
         else :
-            assert 'rcut' in self.descrpt_param, "Error: descriptor must have attr rcut!"
             self.descrpt.enable_compression(self.model_param['compress']["min_nbor_dist"], self.model_param['compress']['model_file'], self.model_param['compress']['table_config'][0], self.model_param['compress']['table_config'][1], self.model_param['compress']['table_config'][2], self.model_param['compress']['table_config'][3])
             self.fitting.init_variables(get_fitting_net_variables(self.model_param['compress']['model_file']))
 

deepmd/utils/graph.py

@@ -112,22 +112,24 @@ def get_tensor_by_type(node,
     return tensor
 
 
-def get_embedding_net_nodes_from_graph_def(graph_def: tf.GraphDef) -> Dict:
+def get_embedding_net_nodes_from_graph_def(graph_def: tf.GraphDef, suffix: str = "") -> Dict:
     """
     Get the embedding net nodes with the given tf.GraphDef object
 
     Parameters
     ----------
     graph_def
         The input tf.GraphDef object
+    suffix : str, optional
+        The scope suffix
     
     Returns
     ----------
     Dict
         The embedding net nodes within the given tf.GraphDef object
     """
     embedding_net_nodes = {}
-    embedding_net_pattern = "filter_type_\d+/matrix_\d+_\d+|filter_type_\d+/bias_\d+_\d+|filter_type_\d+/idt_\d+_\d+|filter_type_all/matrix_\d+_\d+|filter_type_all/bias_\d+_\d+|filter_type_all/idt_\d+_\d"
+    embedding_net_pattern = f"filter_type_\d+{suffix}/matrix_\d+_\d+|filter_type_\d+{suffix}/bias_\d+_\d+|filter_type_\d+{suffix}/idt_\d+_\d+|filter_type_all{suffix}/matrix_\d+_\d+|filter_type_all{suffix}/bias_\d+_\d+|filter_type_all{suffix}/idt_\d+_\d"
     for node in graph_def.node:
         if re.fullmatch(embedding_net_pattern, node.name) != None:
             embedding_net_nodes[node.name] = node.attr["value"].tensor

deepmd/utils/tabulate.py

@@ -34,29 +34,34 @@ class DPTabulate():
             For example, `[[0, 1]]` means no interaction between type 0 and type 1.
     activation_function
             The activation function in the embedding net. Supported options are {"tanh","gelu"} in common.ACTIVATION_FN_DICT.
+    suffix : str, optional
+            The suffix of the scope
     """
     def __init__(self,
                  model_file : str,
                  type_one_side : bool = False,
                  exclude_types : List[List[int]] = [],
-                 activation_fn : Callable[[tf.Tensor], tf.Tensor] = tf.nn.tanh) -> None:
+                 activation_fn : Callable[[tf.Tensor], tf.Tensor] = tf.nn.tanh,
+                 suffix : str = "",
+                 ) -> None:
         """
         Constructor
         """
 
         self.model_file = model_file
         self.type_one_side = type_one_side
         self.exclude_types = exclude_types
+        self.suffix = suffix
         if self.type_one_side and len(self.exclude_types) != 0:
-            raise RunTimeError('"type_one_side" is not compatible with "exclude_types"')
+            raise RuntimeError('"type_one_side" is not compatible with "exclude_types"')
 
         # functype
         if activation_fn == ACTIVATION_FN_DICT["tanh"]:
             self.functype = 1
         elif activation_fn == ACTIVATION_FN_DICT["gelu"]:
             self.functype = 2
         else:
-            raise RunTimeError("Unknown actication function type!")
+            raise RuntimeError("Unknown actication function type!")
         self.activation_fn = activation_fn
 
         self.graph, self.graph_def = load_graph_def(self.model_file)
@@ -72,15 +77,15 @@ def __init__(self,
             self.sel_a = self.graph.get_operation_by_name('DescrptSeA').get_attr('sel_a')
             self.descrpt = self.graph.get_operation_by_name ('DescrptSeA')
 
-        self.davg = get_tensor_by_name_from_graph(self.graph, 'descrpt_attr/t_avg')
-        self.dstd = get_tensor_by_name_from_graph(self.graph, 'descrpt_attr/t_std')
+        self.davg = get_tensor_by_name_from_graph(self.graph, f'descrpt_attr{self.suffix}/t_avg')
+        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.rcut = self.descrpt.get_attr('rcut_r')
         self.rcut_smth = self.descrpt.get_attr('rcut_r_smth')
 
-        self.embedding_net_nodes = get_embedding_net_nodes_from_graph_def(self.graph_def)
+        self.embedding_net_nodes = get_embedding_net_nodes_from_graph_def(self.graph_def, suffix=self.suffix)
 
         for tt in self.exclude_types:
             if (tt[0] not in range(self.ntypes)) or (tt[1] not in range(self.ntypes)):
@@ -174,14 +179,14 @@ def _get_bias(self):
             bias["layer_" + str(layer)] = []
             if self.type_one_side:
                 for ii in range(0, self.ntypes):
-                    tensor_value = np.frombuffer (self.embedding_net_nodes["filter_type_all/bias_" + str(layer) + "_" + str(ii)].tensor_content)
-                    tensor_shape = tf.TensorShape(self.embedding_net_nodes["filter_type_all/bias_" + str(layer) + "_" + str(ii)].tensor_shape).as_list()
+                    tensor_value = np.frombuffer (self.embedding_net_nodes[f"filter_type_all{self.suffix}/bias_{layer}_{ii}"].tensor_content)
+                    tensor_shape = tf.TensorShape(self.embedding_net_nodes[f"filter_type_all{self.suffix}/bias_{layer}_{ii}"].tensor_shape).as_list()
                     bias["layer_" + str(layer)].append(np.reshape(tensor_value, tensor_shape))
             else:
                 for ii in range(0, self.ntypes * self.ntypes):
                     if (ii // self.ntypes, int(ii % self.ntypes)) not in self.exclude_types:
-                        tensor_value = np.frombuffer(self.embedding_net_nodes["filter_type_" + str(ii // self.ntypes) + "/bias_" + str(layer) + "_" + str(int(ii % self.ntypes))].tensor_content)
-                        tensor_shape = tf.TensorShape(self.embedding_net_nodes["filter_type_" + str(ii // self.ntypes) + "/bias_" + str(layer) + "_" + str(int(ii % self.ntypes))].tensor_shape).as_list()
+                        tensor_value = np.frombuffer(self.embedding_net_nodes[f"filter_type_{ii // self.ntypes}{self.suffix}/bias_{layer}_{ii % self.ntypes}"].tensor_content)
+                        tensor_shape = tf.TensorShape(self.embedding_net_nodes[f"filter_type_{ii // self.ntypes}{self.suffix}/bias_{layer}_{ii % self.ntypes}"].tensor_shape).as_list()
                         bias["layer_" + str(layer)].append(np.reshape(tensor_value, tensor_shape))
                     else:
                         bias["layer_" + str(layer)].append(np.array([]))
@@ -193,14 +198,14 @@ def _get_matrix(self):
             matrix["layer_" + str(layer)] = []
             if self.type_one_side:
                 for ii in range(0, self.ntypes):
-                    tensor_value = np.frombuffer (self.embedding_net_nodes["filter_type_all/matrix_" + str(layer) + "_" + str(ii)].tensor_content)
-                    tensor_shape = tf.TensorShape(self.embedding_net_nodes["filter_type_all/matrix_" + str(layer) + "_" + str(ii)].tensor_shape).as_list()
+                    tensor_value = np.frombuffer (self.embedding_net_nodes[f"filter_type_all{self.suffix}/matrix_{layer}_{ii}"].tensor_content)
+                    tensor_shape = tf.TensorShape(self.embedding_net_nodes[f"filter_type_all{self.suffix}/matrix_{layer}_{ii}"].tensor_shape).as_list()
                     matrix["layer_" + str(layer)].append(np.reshape(tensor_value, tensor_shape))
             else:
                 for ii in range(0, self.ntypes * self.ntypes):
                     if (ii // self.ntypes, int(ii % self.ntypes)) not in self.exclude_types:
-                        tensor_value = np.frombuffer(self.embedding_net_nodes["filter_type_" + str(ii // self.ntypes) + "/matrix_" + str(layer) + "_" + str(int(ii % self.ntypes))].tensor_content)
-                        tensor_shape = tf.TensorShape(self.embedding_net_nodes["filter_type_" + str(ii // self.ntypes) + "/matrix_" + str(layer) + "_" + str(int(ii % self.ntypes))].tensor_shape).as_list()
+                        tensor_value = np.frombuffer(self.embedding_net_nodes[f"filter_type_{ii // self.ntypes}{self.suffix}/matrix_{layer}_{ii % self.ntypes}"].tensor_content)
+                        tensor_shape = tf.TensorShape(self.embedding_net_nodes[f"filter_type_{ii // self.ntypes}{self.suffix}/matrix_{layer}_{ii % self.ntypes}"].tensor_shape).as_list()
                         matrix["layer_" + str(layer)].append(np.reshape(tensor_value, tensor_shape))
                     else:
                         matrix["layer_" + str(layer)].append(np.array([]))