Model compression for se-t descriptor

deepmd/descriptor/se_a.py

@@ -323,12 +323,19 @@ def enable_compression(self,
         suffix : str, optional
                 The suffix of the scope
         """
+        # do some checks before the mocel compression process
         assert (
             not self.filter_resnet_dt
         ), "Model compression error: descriptor resnet_dt must be false!"
+        for tt in self.exclude_types:
+            if (tt[0] not in range(self.ntypes)) or (tt[1] not in range(self.ntypes)):
+                raise RuntimeError("exclude types" + str(tt) + " must within the number of atomic types " + str(self.ntypes) + "!")
+        if (self.ntypes * self.ntypes - len(self.exclude_types) == 0):
+            raise RuntimeError("empty embedding-net are not supported in model compression!")
+
         self.compress = True
         self.table = DPTabulate(
-            model_file, self.type_one_side, self.exclude_types, self.compress_activation_fn, suffix=suffix)
+            self, self.filter_neuron, 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, 
@@ -686,7 +693,7 @@ def _filter_lower(
             net = 'filter_-1_net_' + str(type_i)
           else:
             net = 'filter_' + str(type_input) + '_net_' + str(type_i)
-          return op_module.tabulate_fusion(tf.cast(self.table.data[net], self.filter_precision), info, xyz_scatter, tf.reshape(inputs_i, [natom, shape_i[1]//4, 4]), last_layer_size = outputs_size[-1])  
+          return op_module.tabulate_fusion_se_a(tf.cast(self.table.data[net], self.filter_precision), info, xyz_scatter, tf.reshape(inputs_i, [natom, shape_i[1]//4, 4]), last_layer_size = outputs_size[-1])  
         else:
           if (not is_exclude):
               xyz_scatter = embedding_net(

deepmd/descriptor/se_t.py

@@ -10,6 +10,8 @@
 from deepmd.env import default_tf_session_config
 from deepmd.utils.network import embedding_net, embedding_net_rand_seed_shift
 from deepmd.utils.sess import run_sess
+from deepmd.utils.graph import load_graph_def, get_tensor_by_name_from_graph
+from deepmd.utils.tabulate import DPTabulate
 from .descriptor import Descriptor
 from .se import DescrptSe
 
@@ -98,6 +100,7 @@ def __init__ (self,
         self.useBN = False
         self.dstd = None
         self.davg = None
+        self.compress = False
         self.embedding_net_variables = None
 
         self.place_holders = {}
@@ -224,6 +227,53 @@ def compute_input_stats (self,
         self.dstd = np.array(all_dstd)
 
 
+    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
+                The nearest distance between atoms
+        model_file
+                The original frozen model, which will be compressed by the program
+        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
+        """
+        assert (
+            not self.filter_resnet_dt
+        ), "Model compression error: descriptor resnet_dt must be false!"
+        self.compress = True
+        self.table = DPTabulate(
+            self, self.filter_neuron, model_file, activation_fn = self.filter_activation_fn, suffix=suffix)
+        self.table_config = [table_extrapolate, table_stride_1 * 10, table_stride_2 * 10, check_frequency]
+        self.lower, self.upper \
+            = self.table.build(min_nbor_dist, 
+                               table_extrapolate, 
+                               table_stride_1 * 10, 
+                               table_stride_2 * 10)
+
+        graph, _ = load_graph_def(model_file)
+        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)
+
+
     def build (self, 
                coord_ : tf.Tensor, 
                atype_ : tf.Tensor,
@@ -497,25 +547,30 @@ def _filter(self,
                     env_ij = tf.einsum('ijm,ikm->ijk', env_i, env_j)
                     # with (natom x nei_type_i x nei_type_j)
                     ebd_env_ij = tf.reshape(env_ij, [-1, 1])
-                    # with (natom x nei_type_i x nei_type_j) x out_size
-                    ebd_env_ij = embedding_net(ebd_env_ij, 
-                                               self.filter_neuron, 
-                                               self.filter_precision, 
-                                               activation_fn = activation_fn, 
-                                               resnet_dt = self.filter_resnet_dt,
-                                               name_suffix = f"_{type_i}_{type_j}",
-                                               stddev = stddev,
-                                               bavg = bavg,
-                                               seed = self.seed,
-                                               trainable = trainable, 
-                                               uniform_seed = self.uniform_seed,
-                                               initial_variables = self.embedding_net_variables,
-                                               )
-                    if (not self.uniform_seed) and (self.seed is not None): self.seed += self.seed_shift
-                    # with natom x nei_type_i x nei_type_j x out_size
-                    ebd_env_ij = tf.reshape(ebd_env_ij, [-1, nei_type_i, nei_type_j, outputs_size[-1]])
-                    # with natom x out_size
-                    res_ij = tf.einsum('ijk,ijkm->im', env_ij, ebd_env_ij)
+                    if self.compress:
+                        info = [self.lower, self.upper, self.upper * self.table_config[0], self.table_config[1], self.table_config[2], self.table_config[3]]
+                        net = 'filter_' + str(type_i) + '_net_' + str(type_j)
+                        res_ij = op_module.tabulate_fusion_se_t(tf.cast(self.table.data[net], self.filter_precision), info, ebd_env_ij, env_ij, last_layer_size = outputs_size[-1]) 
+                    else:
+                        # with (natom x nei_type_i x nei_type_j) x out_size
+                        ebd_env_ij = embedding_net(ebd_env_ij, 
+                                                   self.filter_neuron, 
+                                                   self.filter_precision, 
+                                                   activation_fn = activation_fn, 
+                                                   resnet_dt = self.filter_resnet_dt,
+                                                   name_suffix = f"_{type_i}_{type_j}",
+                                                   stddev = stddev,
+                                                   bavg = bavg,
+                                                   seed = self.seed,
+                                                   trainable = trainable, 
+                                                   uniform_seed = self.uniform_seed,
+                                                   initial_variables = self.embedding_net_variables,
+                                                   )
+                        if (not self.uniform_seed) and (self.seed is not None): self.seed += self.seed_shift
+                        # with natom x nei_type_i x nei_type_j x out_size
+                        ebd_env_ij = tf.reshape(ebd_env_ij, [-1, nei_type_i, nei_type_j, outputs_size[-1]])
+                        # with natom x out_size
+                        res_ij = tf.einsum('ijk,ijkm->im', env_ij, ebd_env_ij)
                     res_ij = res_ij * (1.0 / float(nei_type_i) / float(nei_type_j))
                     if result is None:
                         result = res_ij

deepmd/entrypoints/compress.py

@@ -93,8 +93,6 @@ def compress(
         name = 'train_attr/min_nbor_dist',
         dtype = GLOBAL_ENER_FLOAT_PRECISION)
     jdata["model"]["compress"] = {}
-    jdata["model"]["compress"]["type"] = 'se_e2_a'
-    jdata["model"]["compress"]["compress"] = True
     jdata["model"]["compress"]["model_file"] = input
     jdata["model"]["compress"]["min_nbor_dist"] = t_min_nbor_dist
     jdata["model"]["compress"]["table_config"] = [

deepmd/entrypoints/convert.py

@@ -1,4 +1,4 @@
-from deepmd.utils.convert import convert_13_to_20, convert_12_to_20
+from deepmd.utils.convert import convert_20_to_21, convert_13_to_21, convert_12_to_21 
 
 def convert(
     *,
@@ -8,8 +8,10 @@ def convert(
     **kwargs,
 ):
     if FROM == '1.2':
-        convert_12_to_20(input_model, output_model)
+        convert_12_to_21(input_model, output_model)
     elif FROM == '1.3':
-        convert_13_to_20(input_model, output_model)
+        convert_13_to_21(input_model, output_model)
+    elif FROM == '2.0':
+        convert_20_to_21(input_model, output_model)
     else:
         raise RuntimeError('unsupported model version ' + FROM)

deepmd/entrypoints/main.py

@@ -386,7 +386,7 @@ def parse_args(args: Optional[List[str]] = None):
     parser_transform.add_argument(
         'FROM',
         type = str,
-        choices = ['1.2', '1.3'],
+        choices = ['1.2', '1.3', '2.0'],
         help="The original model compatibility",
     )
     parser_transform.add_argument(

deepmd/utils/argcheck.py

@@ -371,13 +371,11 @@ def modifier_variant_type_args():
 
 #  --- model compression configurations: --- #
 def model_compression():
-    doc_compress = f"The name of the frozen model file."
     doc_model_file = f"The input model file, which will be compressed by the DeePMD-kit."
     doc_table_config = f"The arguments of model compression, including extrapolate(scale of model extrapolation), stride(uniform stride of tabulation's first and second table), and frequency(frequency of tabulation overflow check)."
     doc_min_nbor_dist = f"The nearest distance between neighbor atoms saved in the frozen model."
 
     return [
-        Argument("compress", bool, optional = False, doc = doc_compress),
         Argument("model_file", str, optional = False, doc = doc_model_file),
         Argument("table_config", list, optional = False, doc = doc_table_config),
         Argument("min_nbor_dist", float, optional = False, doc = doc_min_nbor_dist),

deepmd/utils/convert.py

@@ -3,22 +3,32 @@
 from google.protobuf import text_format
 from tensorflow.python.platform import gfile
 
-def convert_13_to_20(input_model: str, output_model: str):
+def convert_13_to_21(input_model: str, output_model: str):
     convert_pb_to_pbtxt(input_model, 'frozen_model.pbtxt')
     convert_dp13_to_dp20('frozen_model.pbtxt')
+    convert_dp20_to_dp21('frozen_model.pbtxt')
     convert_pbtxt_to_pb('frozen_model.pbtxt', output_model)
     if os.path.isfile('frozen_model.pbtxt'):
         os.remove('frozen_model.pbtxt')
-    print("the converted output model (2.0 support) is saved in %s" % output_model)
+    print("the converted output model (2.1 support) is saved in %s" % output_model)
 
-def convert_12_to_20(input_model: str, output_model: str):
+def convert_12_to_21(input_model: str, output_model: str):
     convert_pb_to_pbtxt(input_model, 'frozen_model.pbtxt')
     convert_dp12_to_dp13('frozen_model.pbtxt')
     convert_dp13_to_dp20('frozen_model.pbtxt')
+    convert_dp20_to_dp21('frozen_model.pbtxt')
     convert_pbtxt_to_pb('frozen_model.pbtxt', output_model)
     if os.path.isfile('frozen_model.pbtxt'):
         os.remove('frozen_model.pbtxt')
-    print("the converted output model (2.0 support) is saved in %s" % output_model)
+    print("the converted output model (2.1 support) is saved in %s" % output_model)
+
+def convert_20_to_21(input_model: str, output_model: str):
+    convert_pb_to_pbtxt(input_model, 'frozen_model.pbtxt')
+    convert_dp20_to_dp21('frozen_model.pbtxt')
+    convert_pbtxt_to_pb('frozen_model.pbtxt', output_model)
+    if os.path.isfile('frozen_model.pbtxt'):
+        os.remove('frozen_model.pbtxt')
+    print("the converted output model (2.1 support) is saved in %s" % output_model)
 
 def convert_pb_to_pbtxt(pbfile: str, pbtxtfile: str):
     with gfile.FastGFile(pbfile, 'rb') as f:
@@ -88,3 +98,60 @@ def convert_dp13_to_dp20(fname: str):
                    .replace('DescrptSeR', 'ProdEnvMatR')
     with open(fname, 'w') as fp:
         fp.write(file_content)
+
+def convert_dp20_to_dp21(fname: str):
+    with open(fname) as fp:
+        file_content = fp.read()
+    old_model_version_node = """
+node {
+  name: "model_attr/model_version"
+  op: "Const"
+  attr {
+    key: "dtype"
+    value {
+      type: DT_STRING
+    }
+  }
+  attr {
+    key: "value"
+    value {
+      tensor {
+        dtype: DT_STRING
+        tensor_shape {
+        }
+        string_val: "1.0"
+      }
+    }
+  }
+}
+"""
+    new_model_version_node = """
+node {
+  name: "model_attr/model_version"
+  op: "Const"
+  attr {
+    key: "dtype"
+    value {
+      type: DT_STRING
+    }
+  }
+  attr {
+    key: "value"
+    value {
+      tensor {
+        dtype: DT_STRING
+        tensor_shape {
+        }
+        string_val: "1.1"
+      }
+    }
+  }
+}
+"""
+    file_content = file_content\
+                   .replace(old_model_version_node, new_model_version_node)\
+                   .replace('TabulateFusion', 'TabulateFusionSeA')\
+                   .replace('TabulateFusionGrad', 'TabulateFusionSeAGrad')\
+                   .replace('TabulateFusionGradGrad', 'TabulateFusionSeAGradGrad')
+    with open(fname, 'w') as fp:
+        fp.write(file_content)