diff --git a/deepmd/descriptor/se_a.py b/deepmd/descriptor/se_a.py index 10ea49abd9..a95f0ca39a 100644 --- a/deepmd/descriptor/se_a.py +++ b/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, diff --git a/deepmd/utils/tabulate.py b/deepmd/utils/tabulate.py index 757189c16d..944428bf48 100644 --- a/deepmd/utils/tabulate.py +++ b/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): diff --git a/source/tests/test_model_compression.py b/source/tests/test_model_compression.py index 91a073c4c2..d67e209c42 100644 --- a/source/tests/test_model_compression.py +++ b/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) \ No newline at end of file