add model compression support for models with exclude_types feature

deepmd/descriptor/se_a.py

@@ -268,7 +268,7 @@ def enable_compression(self,
         self.compress = True
         self.model_file = model_file
         self.table_config = [table_extrapolate, table_stride_1, table_stride_2, check_frequency]
-        self.table = DeepTabulate(self.model_file, self.type_one_side)
+        self.table = DeepTabulate(self.model_file, self.type_one_side, self.exclude_types)
         self.lower, self.upper \
             = self.table.build(min_nbor_dist, 
                                table_extrapolate, 

deepmd/utils/tabulate.py

@@ -19,7 +19,8 @@ class DeepTabulate():
     """
     def __init__(self,
                  model_file : str,
-                 type_one_side : bool = False) -> None:
+                 type_one_side : bool = False,
+                 exclude_types : List[List[int]] = []) -> None:
         """
         Constructor
 
@@ -29,10 +30,15 @@ def __init__(self,
                 The frozen model
         type_one_side
                 Try to build N_types tables. Otherwise, building N_types^2 tables
+        exclude_types : list[list[int]]
+                The Excluded types
         """
 
         self.model_file = model_file
         self.type_one_side = type_one_side
+        self.exclude_types = exclude_types
+        if self.type_one_side and len(self.exclude_types) != 0:
+            raise RunTimeError('"type_one_side" is not compatible with "exclude_types"')
 
         self.graph, self.graph_def = self._load_graph()
         self.sess = tf.Session(graph = self.graph)
@@ -55,19 +61,26 @@ def __init__(self,
         self.rcut_smth = self.descrpt.get_attr('rcut_r_smth')
 
         self.filter_variable_nodes = self._load_matrix_node()
-        self.layer_size = int(len(self.filter_variable_nodes) / (self.ntypes * self.ntypes * 2))
+        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.layer_size = len(self.filter_variable_nodes) // ((self.ntypes * self.ntypes - len(self.exclude_types)) * 2)
         self.table_size = self.ntypes * self.ntypes
         if type_one_side :
-            self.layer_size = int(len(self.filter_variable_nodes) / (self.ntypes * 2))
+            self.layer_size = len(self.filter_variable_nodes) // (self.ntypes * 2)
             self.table_size = self.ntypes
         # self.value_type = self.filter_variable_nodes["filter_type_0/matrix_1_0"].dtype #"filter_type_0/matrix_1_0" must exit~
         # get trained variables
         self.bias = self._get_bias()
         self.matrix = self._get_matrix()
 
-        self.data_type = type(self.matrix["layer_1"][0][0][0])
-        assert self.matrix["layer_1"][0].size > 0, "no matrix exist in matrix array!"
-        self.last_layer_size = self.matrix["layer_" + str(self.layer_size)][0].shape[1]
+        for item in self.matrix["layer_" + str(self.layer_size)]:
+            if len(item) != 0:
+                self.data_type = type(item[0][0])
+                self.last_layer_size = item.shape[1]
         # define tables
         self.data = {}
 
@@ -91,7 +104,7 @@ def build(self,
                 The uniform stride of the first table
         stride1
                 The uniform stride of the second table
-        
+
         Returns
         ----------
         lower
@@ -106,27 +119,27 @@ def build(self,
         xx = np.append(xx, np.array([extrapolate * upper], dtype = self.data_type))
         self.nspline = int((upper - lower) / stride0 + (extrapolate * upper - upper) / stride1)
         for ii in range(self.table_size):
-            vv, dd, d2 = self._make_data(xx, ii)
-            if self.type_one_side:
-                net = "filter_-1_net_" + str(int(ii))
-            else:
-                net = "filter_" + str(int(ii / self.ntypes)) + "_net_" + str(int(ii % self.ntypes))
-            self.data[net] = np.zeros([self.nspline, 6 * self.last_layer_size], dtype = self.data_type)
-            # for jj in tqdm(range(self.nspline), desc = 'DEEPMD INFO    |-> deepmd.utils.tabulate\t\t\t' + net + ', tabulating'):
-            for jj in range(self.nspline):
-                for kk in range(self.last_layer_size):
-                    if jj < int((upper - lower) / stride0):
-                        tt = stride0
-                    else:
-                        tt = stride1
-                    hh = vv[jj + 1][kk] - vv[jj][kk]
-                    self.data[net][jj][kk * 6 + 0] = vv[jj][kk]
-                    self.data[net][jj][kk * 6 + 1] = dd[jj][kk]
-                    self.data[net][jj][kk * 6 + 2] = 0.5 * d2[jj][kk]
-                    self.data[net][jj][kk * 6 + 3] = (1 / (2 * tt * tt * tt)) * (20 * hh - (8 * dd[jj + 1][kk] + 12 * dd[jj][kk]) * tt - (3 * d2[jj][kk] - d2[jj + 1][kk]) * tt * tt)
-                    self.data[net][jj][kk * 6 + 4] = (1 / (2 * tt * tt * tt * tt)) * (-30 * hh + (14 * dd[jj + 1][kk] + 16 * dd[jj][kk]) * tt + (3 * d2[jj][kk] - 2 * d2[jj + 1][kk]) * tt * tt)
-                    self.data[net][jj][kk * 6 + 5] = (1 / (2 * tt * tt * tt * tt * tt)) * (12 * hh - 6 * (dd[jj + 1][kk] + dd[jj][kk]) * tt + (d2[jj + 1][kk] - d2[jj][kk]) * tt * tt)
-            self.data[net]
+            if self.type_one_side or (ii // self.ntypes, int(ii % self.ntypes)) not in self.exclude_types:
+                vv, dd, d2 = self._make_data(xx, ii)
+                if self.type_one_side:
+                    net = "filter_-1_net_" + str(ii)
+                else:
+                    net = "filter_" + str(ii // self.ntypes) + "_net_" + str(int(ii % self.ntypes))
+                self.data[net] = np.zeros([self.nspline, 6 * self.last_layer_size], dtype = self.data_type)
+                # for jj in tqdm(range(self.nspline), desc = 'DEEPMD INFO    |-> deepmd.utils.tabulate\t\t\t' + net + ', tabulating'):
+                for jj in range(self.nspline):
+                    for kk in range(self.last_layer_size):
+                        if jj < int((upper - lower) / stride0):
+                            tt = stride0
+                        else:
+                            tt = stride1
+                        hh = vv[jj + 1][kk] - vv[jj][kk]
+                        self.data[net][jj][kk * 6 + 0] = vv[jj][kk]
+                        self.data[net][jj][kk * 6 + 1] = dd[jj][kk]
+                        self.data[net][jj][kk * 6 + 2] = 0.5 * d2[jj][kk]
+                        self.data[net][jj][kk * 6 + 3] = (1 / (2 * tt * tt * tt)) * (20 * hh - (8 * dd[jj + 1][kk] + 12 * dd[jj][kk]) * tt - (3 * d2[jj][kk] - d2[jj + 1][kk]) * tt * tt)
+                        self.data[net][jj][kk * 6 + 4] = (1 / (2 * tt * tt * tt * tt)) * (-30 * hh + (14 * dd[jj + 1][kk] + 16 * dd[jj][kk]) * tt + (3 * d2[jj][kk] - 2 * d2[jj + 1][kk]) * tt * tt)
+                        self.data[net][jj][kk * 6 + 5] = (1 / (2 * tt * tt * tt * tt * tt)) * (12 * hh - 6 * (dd[jj + 1][kk] + dd[jj][kk]) * tt + (d2[jj + 1][kk] - d2[jj][kk]) * tt * tt)
         return lower, upper
 
     def _load_graph(self):
@@ -165,14 +178,17 @@ 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.filter_variable_nodes["filter_type_all/bias_" + str(layer) + "_" + str(int(ii))].tensor_content)
-                    tensor_shape = tf.TensorShape(self.filter_variable_nodes["filter_type_all/bias_" + str(layer) + "_" + str(int(ii))].tensor_shape).as_list()
+                    tensor_value = np.frombuffer (self.filter_variable_nodes["filter_type_all/bias_" + str(layer) + "_" + str(ii)].tensor_content)
+                    tensor_shape = tf.TensorShape(self.filter_variable_nodes["filter_type_all/bias_" + str(layer) + "_" + str(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):
-                    tensor_value = np.frombuffer(self.filter_variable_nodes["filter_type_" + str(int(ii / self.ntypes)) + "/bias_" + str(layer) + "_" + str(int(ii % self.ntypes))].tensor_content)
-                    tensor_shape = tf.TensorShape(self.filter_variable_nodes["filter_type_" + str(int(ii / self.ntypes)) + "/bias_" + str(layer) + "_" + str(int(ii % self.ntypes))].tensor_shape).as_list()
-                    bias["layer_" + str(layer)].append(np.reshape(tensor_value, tensor_shape))
+                    if (ii // self.ntypes, int(ii % self.ntypes)) not in self.exclude_types:
+                        tensor_value = np.frombuffer(self.filter_variable_nodes["filter_type_" + str(ii // self.ntypes) + "/bias_" + str(layer) + "_" + str(int(ii % self.ntypes))].tensor_content)
+                        tensor_shape = tf.TensorShape(self.filter_variable_nodes["filter_type_" + str(ii // self.ntypes) + "/bias_" + str(layer) + "_" + str(int(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([]))
         return bias
 
     def _get_matrix(self):
@@ -181,14 +197,17 @@ 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.filter_variable_nodes["filter_type_all/matrix_" + str(layer) + "_" + str(int(ii))].tensor_content)
-                    tensor_shape = tf.TensorShape(self.filter_variable_nodes["filter_type_all/matrix_" + str(layer) + "_" + str(int(ii))].tensor_shape).as_list()
+                    tensor_value = np.frombuffer (self.filter_variable_nodes["filter_type_all/matrix_" + str(layer) + "_" + str(ii)].tensor_content)
+                    tensor_shape = tf.TensorShape(self.filter_variable_nodes["filter_type_all/matrix_" + str(layer) + "_" + str(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):
-                    tensor_value = np.frombuffer(self.filter_variable_nodes["filter_type_" + str(int(ii / self.ntypes)) + "/matrix_" + str(layer) + "_" + str(int(ii % self.ntypes))].tensor_content)
-                    tensor_shape = tf.TensorShape(self.filter_variable_nodes["filter_type_" + str(int(ii / self.ntypes)) + "/matrix_" + str(layer) + "_" + str(int(ii % self.ntypes))].tensor_shape).as_list()
-                    matrix["layer_" + str(layer)].append(np.reshape(tensor_value, tensor_shape))
+                    if (ii // self.ntypes, int(ii % self.ntypes)) not in self.exclude_types:
+                        tensor_value = np.frombuffer(self.filter_variable_nodes["filter_type_" + str(ii // self.ntypes) + "/matrix_" + str(layer) + "_" + str(int(ii % self.ntypes))].tensor_content)
+                        tensor_shape = tf.TensorShape(self.filter_variable_nodes["filter_type_" + str(ii // self.ntypes) + "/matrix_" + str(layer) + "_" + str(int(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([]))
         return matrix
 
     # one-by-one executions
@@ -221,8 +240,8 @@ def _layer_1(self, x, w, b):
 
     def _save_data(self):
         for ii in range(self.ntypes * self.ntypes):
-            net = "filter_" + str(int(ii / self.ntypes)) + "_net_" + str(int(ii % self.ntypes))
-            np.savetxt('data_' + str(int(ii)), self.data[net])
+            net = "filter_" + str(ii // self.ntypes) + "_net_" + str(int(ii % self.ntypes))
+            np.savetxt('data_' + str(ii), self.data[net])
 
     def _get_env_mat_range(self,
                            min_nbor_dist):

source/tests/test_model_compression.py

@@ -384,3 +384,137 @@ def test_ase(self):
         for ii in range(nframes):
             self.assertAlmostEqual(ee0.reshape([-1])[ii], ee1.reshape([-1])[ii], places = default_places)
 
+class TestDeepPotAPBCExcludeTypes(unittest.TestCase) :
+    def setUp(self):
+        self.data_file  = str(tests_path / os.path.join("model_compression", "data"))
+        self.frozen_model = str(tests_path / "dp-original.pb")
+        self.compressed_model = str(tests_path / "dp-compressed.pb")
+        self.INPUT = str(tests_path / "input.json")
+        jdata = j_loader(str(tests_path / os.path.join("model_compression", "input.json")))
+        jdata["training"]["training_data"]["systems"] = self.data_file
+        jdata["training"]["validation_data"]["systems"] = self.data_file
+        jdata["model"]["descriptor"]["exclude_types"] = [[0, 1]]
+        with open(self.INPUT, "w") as fp:
+            json.dump(jdata, fp, indent=4)
+
+        ret = os.system("dp train " + self.INPUT)
+        assert(ret == 0), "DP train error!"
+        ret = os.system("dp freeze -o " + self.frozen_model)
+        assert(ret == 0), "DP freeze error!"
+        ret = os.system("dp compress " + self.INPUT + " -i " + self.frozen_model + " -o " + self.compressed_model)
+        assert(ret == 0), "DP model compression error!"
+
+        self.dp_original = DeepPot(self.frozen_model)
+        self.dp_compressed = DeepPot(self.compressed_model)
+        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])
+        self.atype = [0, 1, 1, 0, 1, 1]
+        self.box = np.array([13., 0., 0., 0., 13., 0., 0., 0., 13.])
+
+    def tearDown(self):
+        _file_delete(self.INPUT)
+        _file_delete(self.frozen_model)
+        _file_delete(self.compressed_model)
+        _file_delete("out.json")
+        _file_delete("compress.json")
+        _file_delete("checkpoint")
+        _file_delete("lcurve.out")
+        _file_delete("model.ckpt.meta")
+        _file_delete("model.ckpt.index")
+        _file_delete("model.ckpt.data-00000-of-00001")
+
+    def test_attrs(self):
+        self.assertEqual(self.dp_original.get_ntypes(), 2)
+        self.assertAlmostEqual(self.dp_original.get_rcut(), 6.0, places = default_places)
+        self.assertEqual(self.dp_original.get_type_map(), ['O', 'H'])
+        self.assertEqual(self.dp_original.get_dim_fparam(), 0)
+        self.assertEqual(self.dp_original.get_dim_aparam(), 0)
+
+        self.assertEqual(self.dp_compressed.get_ntypes(), 2)
+        self.assertAlmostEqual(self.dp_compressed.get_rcut(), 6.0, places = default_places)
+        self.assertEqual(self.dp_compressed.get_type_map(), ['O', 'H'])
+        self.assertEqual(self.dp_compressed.get_dim_fparam(), 0)
+        self.assertEqual(self.dp_compressed.get_dim_aparam(), 0)
+
+    def test_1frame(self):
+        ee0, ff0, vv0 = self.dp_original.eval(self.coords, self.box, self.atype, atomic = False)
+        ee1, ff1, vv1 = self.dp_compressed.eval(self.coords, self.box, self.atype, atomic = False)
+        # check shape of the returns
+        nframes = 1
+        natoms = len(self.atype)
+        self.assertEqual(ee0.shape, (nframes,1))
+        self.assertEqual(ff0.shape, (nframes,natoms,3))
+        self.assertEqual(vv0.shape, (nframes,9))
+        self.assertEqual(ee1.shape, (nframes,1))
+        self.assertEqual(ff1.shape, (nframes,natoms,3))
+        self.assertEqual(vv1.shape, (nframes,9))
+        # check values
+        for ii in range(ff0.size):
+            self.assertAlmostEqual(ff0.reshape([-1])[ii], ff1.reshape([-1])[ii], places = default_places)
+        for ii in range(nframes):
+            self.assertAlmostEqual(ee0.reshape([-1])[ii], ee1.reshape([-1])[ii], places = default_places)
+        for ii in range(nframes, 9):
+            self.assertAlmostEqual(vv0.reshape([-1])[ii], vv1.reshape([-1])[ii], places = default_places)
+
+    def test_1frame_atm(self):
+        ee0, ff0, vv0, ae0, av0 = self.dp_original.eval(self.coords, self.box, self.atype, atomic = True)
+        ee1, ff1, vv1, ae1, av1 = self.dp_compressed.eval(self.coords, self.box, self.atype, atomic = True)
+        # check shape of the returns
+        nframes = 1
+        natoms = len(self.atype)
+        self.assertEqual(ee0.shape, (nframes,1))
+        self.assertEqual(ff0.shape, (nframes,natoms,3))
+        self.assertEqual(vv0.shape, (nframes,9))
+        self.assertEqual(ae0.shape, (nframes,natoms,1))
+        self.assertEqual(av0.shape, (nframes,natoms,9))
+        self.assertEqual(ee1.shape, (nframes,1))
+        self.assertEqual(ff1.shape, (nframes,natoms,3))
+        self.assertEqual(vv1.shape, (nframes,9))
+        self.assertEqual(ae1.shape, (nframes,natoms,1))
+        self.assertEqual(av1.shape, (nframes,natoms,9))
+        # check values
+        for ii in range(ff0.size):
+            self.assertAlmostEqual(ff0.reshape([-1])[ii], ff1.reshape([-1])[ii], places = default_places)
+        for ii in range(ae0.size):
+            self.assertAlmostEqual(ae0.reshape([-1])[ii], ae1.reshape([-1])[ii], places = default_places)
+        for ii in range(av0.size):
+            self.assertAlmostEqual(av0.reshape([-1])[ii], av1.reshape([-1])[ii], places = default_places)
+        for ii in range(nframes):
+            self.assertAlmostEqual(ee0.reshape([-1])[ii], ee1.reshape([-1])[ii], places = default_places)
+        for ii in range(nframes, 9):
+            self.assertAlmostEqual(vv0.reshape([-1])[ii], vv1.reshape([-1])[ii], places = default_places)
+
+    def test_2frame_atm(self):
+        coords2 = np.concatenate((self.coords, self.coords))
+        box2 = np.concatenate((self.box, self.box))
+        ee0, ff0, vv0, ae0, av0 = self.dp_original.eval(coords2, box2, self.atype, atomic = True)
+        ee1, ff1, vv1, ae1, av1 = self.dp_compressed.eval(coords2, box2, self.atype, atomic = True)
+        # check shape of the returns
+        nframes = 2
+        natoms = len(self.atype)
+        self.assertEqual(ee0.shape, (nframes,1))
+        self.assertEqual(ff0.shape, (nframes,natoms,3))
+        self.assertEqual(vv0.shape, (nframes,9))
+        self.assertEqual(ae0.shape, (nframes,natoms,1))
+        self.assertEqual(av0.shape, (nframes,natoms,9))
+        self.assertEqual(ee1.shape, (nframes,1))
+        self.assertEqual(ff1.shape, (nframes,natoms,3))
+        self.assertEqual(vv1.shape, (nframes,9))
+        self.assertEqual(ae1.shape, (nframes,natoms,1))
+        self.assertEqual(av1.shape, (nframes,natoms,9))
+
+        # check values
+        for ii in range(ff0.size):
+            self.assertAlmostEqual(ff0.reshape([-1])[ii], ff1.reshape([-1])[ii], places = default_places)
+        for ii in range(ae0.size):
+            self.assertAlmostEqual(ae0.reshape([-1])[ii], ae1.reshape([-1])[ii], places = default_places)
+        for ii in range(av0.size):
+            self.assertAlmostEqual(av0.reshape([-1])[ii], av1.reshape([-1])[ii], places = default_places)
+        for ii in range(nframes):
+            self.assertAlmostEqual(ee0.reshape([-1])[ii], ee1.reshape([-1])[ii], places = default_places)
+        for ii in range(nframes, 9):
+            self.assertAlmostEqual(vv0.reshape([-1])[ii], vv1.reshape([-1])[ii], places = default_places)