From a4e01138443170ca0700b209b13de01ccb23506f Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Wed, 12 Jan 2022 09:39:33 -0500 Subject: [PATCH 1/7] support `type_one_side` with `exclude_types` --- deepmd/descriptor/se_a.py | 16 ++++++++++++---- 1 file changed, 12 insertions(+), 4 deletions(-) diff --git a/deepmd/descriptor/se_a.py b/deepmd/descriptor/se_a.py index 9883909082..c2e8d931d6 100644 --- a/deepmd/descriptor/se_a.py +++ b/deepmd/descriptor/se_a.py @@ -142,8 +142,6 @@ def __init__ (self, self.exclude_types.add((tt[1], tt[0])) self.set_davg_zero = set_davg_zero self.type_one_side = type_one_side - if self.type_one_side and len(exclude_types) != 0: - raise RuntimeError('"type_one_side" is not compatible with "exclude_types"') # descrpt config self.sel_r = [ 0 for ii in range(len(self.sel_a)) ] @@ -552,13 +550,23 @@ def _pass_filter(self, inputs = tf.reshape(inputs, [-1, self.ndescrpt * natoms[0]]) output = [] output_qmat = [] - if not self.type_one_side and type_embedding is None: + if not (self.type_one_side and len(self.exclude_types) == 0) and type_embedding is None: for type_i in range(self.ntypes): inputs_i = tf.slice (inputs, [ 0, start_index* self.ndescrpt], [-1, natoms[2+type_i]* self.ndescrpt] ) inputs_i = tf.reshape(inputs_i, [-1, self.ndescrpt]) - layer, qmat = self._filter(inputs_i, type_i, name='filter_type_'+str(type_i)+suffix, natoms=natoms, reuse=reuse, trainable = trainable, activation_fn = self.filter_activation_fn) + if self.type_one_side: + # reuse NN parameters for all types, which should be equivalent to lines below `else` + # the number of NN parameters are reduced + # however, we can still apply exclude_types by this way + # Note: compared to lines below `else`, it cannot utilize 100% GPU + # but ops on CPUs have the same performance + reuse = tf.AUTO_REUSE + filter_name = 'filter_type_all'+suffix + else: + filter_name = 'filter_type_'+str(type_i)+suffix + layer, qmat = self._filter(inputs_i, type_i, name=filter_name, natoms=natoms, reuse=reuse, trainable = trainable, activation_fn = self.filter_activation_fn) layer = tf.reshape(layer, [tf.shape(inputs)[0], natoms[2+type_i] * self.get_dim_out()]) qmat = tf.reshape(qmat, [tf.shape(inputs)[0], natoms[2+type_i] * self.get_dim_rot_mat_1() * 3]) output.append(layer) From 10b7baf5978e1e726e28d356e1301f2a1ed9f9d4 Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Wed, 12 Jan 2022 10:07:50 -0500 Subject: [PATCH 2/7] add support for se_r --- deepmd/descriptor/se_r.py | 14 ++++++++++++-- 1 file changed, 12 insertions(+), 2 deletions(-) diff --git a/deepmd/descriptor/se_r.py b/deepmd/descriptor/se_r.py index c92c51225d..fdbf5873d6 100644 --- a/deepmd/descriptor/se_r.py +++ b/deepmd/descriptor/se_r.py @@ -443,13 +443,23 @@ def _pass_filter(self, start_index = 0 inputs = tf.reshape(inputs, [-1, self.ndescrpt * natoms[0]]) output = [] - if not self.type_one_side: + if not (self.type_one_side and len(self.exclude_types) == 0): for type_i in range(self.ntypes): inputs_i = tf.slice (inputs, [ 0, start_index* self.ndescrpt], [-1, natoms[2+type_i]* self.ndescrpt] ) inputs_i = tf.reshape(inputs_i, [-1, self.ndescrpt]) - layer = self._filter_r(inputs_i, type_i, name='filter_type_'+str(type_i)+suffix, natoms=natoms, reuse=reuse, trainable = trainable, activation_fn = self.filter_activation_fn) + if self.type_one_side: + # reuse NN parameters for all types, which should be equivalent to lines below `else` + # the number of NN parameters are reduced + # however, we can still apply exclude_types by this way + # Note: compared to lines below `else`, it cannot utilize 100% GPU + # but ops on CPUs have the same performance + reuse = tf.AUTO_REUSE + filter_name = 'filter_type_all'+suffix + else: + filter_name = 'filter_type_'+str(type_i)+suffix + layer = self._filter_r(inputs_i, type_i, name=filter_name, natoms=natoms, reuse=reuse, trainable = trainable, activation_fn = self.filter_activation_fn) layer = tf.reshape(layer, [tf.shape(inputs)[0], natoms[2+type_i] * self.get_dim_out()]) output.append(layer) start_index += natoms[2+type_i] From 41f62f5aee7d95b0624d3d29ac7dd94f6873e11e Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Wed, 12 Jan 2022 11:15:56 -0500 Subject: [PATCH 3/7] fix compress issues --- deepmd/descriptor/se_a.py | 2 +- deepmd/descriptor/se_r.py | 6 +--- deepmd/utils/tabulate.py | 59 +++++++++++++++++++++++++++++++-------- 3 files changed, 49 insertions(+), 18 deletions(-) diff --git a/deepmd/descriptor/se_a.py b/deepmd/descriptor/se_a.py index c2e8d931d6..3326364700 100644 --- a/deepmd/descriptor/se_a.py +++ b/deepmd/descriptor/se_a.py @@ -560,7 +560,7 @@ def _pass_filter(self, # reuse NN parameters for all types, which should be equivalent to lines below `else` # the number of NN parameters are reduced # however, we can still apply exclude_types by this way - # Note: compared to lines below `else`, it cannot utilize 100% GPU + # Note: compared to lines below `else`, it cannot utilize 100% GPU, as TF doesn't allow multi-GPU streaming # but ops on CPUs have the same performance reuse = tf.AUTO_REUSE filter_name = 'filter_type_all'+suffix diff --git a/deepmd/descriptor/se_r.py b/deepmd/descriptor/se_r.py index fdbf5873d6..90b8cb9fba 100644 --- a/deepmd/descriptor/se_r.py +++ b/deepmd/descriptor/se_r.py @@ -450,11 +450,7 @@ def _pass_filter(self, [-1, natoms[2+type_i]* self.ndescrpt] ) inputs_i = tf.reshape(inputs_i, [-1, self.ndescrpt]) if self.type_one_side: - # reuse NN parameters for all types, which should be equivalent to lines below `else` - # the number of NN parameters are reduced - # however, we can still apply exclude_types by this way - # Note: compared to lines below `else`, it cannot utilize 100% GPU - # but ops on CPUs have the same performance + # see comments in se_a reuse = tf.AUTO_REUSE filter_name = 'filter_type_all'+suffix else: diff --git a/deepmd/utils/tabulate.py b/deepmd/utils/tabulate.py index e3107821a8..662f3f6dcf 100644 --- a/deepmd/utils/tabulate.py +++ b/deepmd/utils/tabulate.py @@ -4,6 +4,7 @@ import deepmd from typing import Callable from typing import Tuple, List +from functools import lru_cache from scipy.special import comb from deepmd.env import tf from deepmd.env import op_module @@ -174,7 +175,7 @@ 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): - if self.type_one_side or (ii // self.ntypes, ii % self.ntypes) not in self.exclude_types: + if (self.type_one_side and not self._all_excluded(ii)) or (not self.type_one_side and (ii // self.ntypes, ii % self.ntypes) not in self.exclude_types): if self.type_one_side: net = "filter_-1_net_" + str(ii) else: @@ -198,7 +199,7 @@ 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): - if self.type_one_side or (ii // self.ntypes, ii % self.ntypes) not in self.exclude_types: + if (self.type_one_side and not self._all_excluded(ii)) or (not self.type_one_side and (ii // self.ntypes, ii % self.ntypes) not in self.exclude_types): if self.type_one_side: net = "filter_-1_net_" + str(ii) else: @@ -249,8 +250,11 @@ def _get_bias(self): if isinstance(self.descrpt, deepmd.descriptor.DescrptSeA): if self.type_one_side: for ii in range(0, self.ntypes): - node = self.embedding_net_nodes[f"filter_type_all{self.suffix}/bias_{layer}_{ii}"] - bias["layer_" + str(layer)].append(tf.make_ndarray(node)) + if not self._all_excluded(ii): + node = self.embedding_net_nodes[f"filter_type_all{self.suffix}/bias_{layer}_{ii}"] + bias["layer_" + str(layer)].append(tf.make_ndarray(node)) + else: + bias["layer_" + str(layer)].append(np.array([])) else: for ii in range(0, self.ntypes * self.ntypes): if (ii // self.ntypes, ii % self.ntypes) not in self.exclude_types: @@ -266,8 +270,11 @@ def _get_bias(self): elif isinstance(self.descrpt, deepmd.descriptor.DescrptSeR): if self.type_one_side: for ii in range(0, self.ntypes): - node = self.embedding_net_nodes[f"filter_type_all{self.suffix}/bias_{layer}_{ii}"] - bias["layer_" + str(layer)].append(tf.make_ndarray(node)) + if not self._all_excluded(ii): + node = self.embedding_net_nodes[f"filter_type_all{self.suffix}/bias_{layer}_{ii}"] + bias["layer_" + str(layer)].append(tf.make_ndarray(node)) + else: + bias["layer_" + str(layer)].append(np.array([])) else: for ii in range(0, self.ntypes * self.ntypes): if (ii // self.ntypes, ii % self.ntypes) not in self.exclude_types: @@ -286,8 +293,11 @@ def _get_matrix(self): if isinstance(self.descrpt, deepmd.descriptor.DescrptSeA): if self.type_one_side: for ii in range(0, self.ntypes): - node = self.embedding_net_nodes[f"filter_type_all{self.suffix}/matrix_{layer}_{ii}"] - matrix["layer_" + str(layer)].append(tf.make_ndarray(node)) + if not self._all_excluded(ii): + node = self.embedding_net_nodes[f"filter_type_all{self.suffix}/matrix_{layer}_{ii}"] + matrix["layer_" + str(layer)].append(tf.make_ndarray(node)) + else: + matrix["layer_" + str(layer)].append(np.array([])) else: for ii in range(0, self.ntypes * self.ntypes): if (ii // self.ntypes, ii % self.ntypes) not in self.exclude_types: @@ -303,8 +313,11 @@ def _get_matrix(self): elif isinstance(self.descrpt, deepmd.descriptor.DescrptSeR): if self.type_one_side: for ii in range(0, self.ntypes): - node = self.embedding_net_nodes[f"filter_type_all{self.suffix}/matrix_{layer}_{ii}"] - matrix["layer_" + str(layer)].append(tf.make_ndarray(node)) + if not self._all_excluded(ii): + node = self.embedding_net_nodes[f"filter_type_all{self.suffix}/matrix_{layer}_{ii}"] + matrix["layer_" + str(layer)].append(tf.make_ndarray(node)) + else: + matrix["layer_" + str(layer)].append(np.array([])) else: for ii in range(0, self.ntypes * self.ntypes): if (ii // self.ntypes, ii % self.ntypes) not in self.exclude_types: @@ -411,16 +424,38 @@ def _get_layer_size(self): if isinstance(self.descrpt, deepmd.descriptor.DescrptSeA): layer_size = len(self.embedding_net_nodes) // ((self.ntypes * self.ntypes - len(self.exclude_types)) * 2) if self.type_one_side : - layer_size = len(self.embedding_net_nodes) // (self.ntypes * 2) + layer_size = len(self.embedding_net_nodes) // ((self.ntypes - self._n_all_excluded) * 2) elif isinstance(self.descrpt, deepmd.descriptor.DescrptSeT): layer_size = len(self.embedding_net_nodes) // int(comb(self.ntypes + 1, 2) * 2) elif isinstance(self.descrpt, deepmd.descriptor.DescrptSeR): layer_size = len(self.embedding_net_nodes) // ((self.ntypes * self.ntypes - len(self.exclude_types)) * 2) if self.type_one_side : - layer_size = len(self.embedding_net_nodes) // (self.ntypes * 2) + layer_size = len(self.embedding_net_nodes) // ((self.ntypes - self._n_all_excluded) * 2) else: raise RuntimeError("Unsupported descriptor") return layer_size + + @property + @lru_cache + def _n_all_excluded(self) -> int: + """Then number of types excluding all types.""" + return sum((int(self._all_excluded(ii)) for ii in range(0, self.ntypes))) + + @lru_cache + def _all_excluded(self, ii: int) -> bool: + """Check if type ii excluds all types. + + Parameters + ---------- + ii : int + type index + + Returns + ------- + bool + if type ii excluds all types + """ + return all([(ii, type_i) in self.exclude_types for type_i in range(self.ntypes)]) def _get_table_size(self): table_size = 0 From 45460c809265c3a3bcc627007cea11f868c2eb9b Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Sun, 16 Jan 2022 05:02:29 -0500 Subject: [PATCH 4/7] rewrite comments --- deepmd/descriptor/se_a.py | 6 +----- deepmd/descriptor/se_r.py | 2 +- 2 files changed, 2 insertions(+), 6 deletions(-) diff --git a/deepmd/descriptor/se_a.py b/deepmd/descriptor/se_a.py index 3326364700..d38adf44c3 100644 --- a/deepmd/descriptor/se_a.py +++ b/deepmd/descriptor/se_a.py @@ -557,11 +557,7 @@ def _pass_filter(self, [-1, natoms[2+type_i]* self.ndescrpt] ) inputs_i = tf.reshape(inputs_i, [-1, self.ndescrpt]) if self.type_one_side: - # reuse NN parameters for all types, which should be equivalent to lines below `else` - # the number of NN parameters are reduced - # however, we can still apply exclude_types by this way - # Note: compared to lines below `else`, it cannot utilize 100% GPU, as TF doesn't allow multi-GPU streaming - # but ops on CPUs have the same performance + # reuse NN parameters for all types to support type_one_side along with exclude_types reuse = tf.AUTO_REUSE filter_name = 'filter_type_all'+suffix else: diff --git a/deepmd/descriptor/se_r.py b/deepmd/descriptor/se_r.py index 90b8cb9fba..b9e5d5aabd 100644 --- a/deepmd/descriptor/se_r.py +++ b/deepmd/descriptor/se_r.py @@ -450,7 +450,7 @@ def _pass_filter(self, [-1, natoms[2+type_i]* self.ndescrpt] ) inputs_i = tf.reshape(inputs_i, [-1, self.ndescrpt]) if self.type_one_side: - # see comments in se_a + # reuse NN parameters for all types to support type_one_side along with exclude_types reuse = tf.AUTO_REUSE filter_name = 'filter_type_all'+suffix else: From dd8f34db8aa4a219104cf4cb499888013579f932 Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Sun, 16 Jan 2022 05:12:50 -0500 Subject: [PATCH 5/7] fix compatibility with Python 3.6 --- deepmd/utils/tabulate.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/deepmd/utils/tabulate.py b/deepmd/utils/tabulate.py index 662f3f6dcf..ee1088bd3c 100644 --- a/deepmd/utils/tabulate.py +++ b/deepmd/utils/tabulate.py @@ -436,12 +436,12 @@ def _get_layer_size(self): return layer_size @property - @lru_cache + @lru_cache() def _n_all_excluded(self) -> int: """Then number of types excluding all types.""" return sum((int(self._all_excluded(ii)) for ii in range(0, self.ntypes))) - @lru_cache + @lru_cache() def _all_excluded(self, ii: int) -> bool: """Check if type ii excluds all types. @@ -483,4 +483,4 @@ def _get_last_layer_size(self): for item in self.matrix["layer_" + str(self.layer_size)]: if len(item) != 0: return item.shape[1] - return 0 \ No newline at end of file + return 0 From f6684a0a1c3bfd786560b76f726e3702c49558d8 Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Sun, 13 Mar 2022 20:24:38 -0400 Subject: [PATCH 6/7] add unittests --- ...ession_se_a_type_one_side_exclude_types.py | 147 ++++++++++++++++++ source/tests/test_type_one_side.py | 135 ++++++++++++++++ 2 files changed, 282 insertions(+) create mode 100644 source/tests/test_model_compression_se_a_type_one_side_exclude_types.py create mode 100644 source/tests/test_type_one_side.py diff --git a/source/tests/test_model_compression_se_a_type_one_side_exclude_types.py b/source/tests/test_model_compression_se_a_type_one_side_exclude_types.py new file mode 100644 index 0000000000..d29fc6b1fe --- /dev/null +++ b/source/tests/test_model_compression_se_a_type_one_side_exclude_types.py @@ -0,0 +1,147 @@ +import os,sys,platform,shutil,dpdata,json +import numpy as np +import unittest +import subprocess as sp + +from deepmd.infer import DeepPot +from deepmd.env import MODEL_VERSION +# from deepmd.entrypoints.compress import compress +from common import j_loader, tests_path + +from deepmd.env import GLOBAL_NP_FLOAT_PRECISION +if GLOBAL_NP_FLOAT_PRECISION == np.float32 : + default_places = 4 +else : + default_places = 10 + +def _file_delete(file) : + if os.path.isdir(file): + os.rmdir(file) + elif os.path.isfile(file): + os.remove(file) + +def _subprocess_run(command): + popen = sp.Popen(command.split(), shell=False, stdout=sp.PIPE, stderr=sp.STDOUT) + for line in iter(popen.stdout.readline, b''): + if hasattr(line, 'decode'): + line = line.decode('utf-8') + line = line.rstrip() + print(line) + popen.wait() + return popen.returncode + +def _init_models(): + data_file = str(tests_path / os.path.join("model_compression", "data")) + frozen_model = str(tests_path / "dp-original.pb") + compressed_model = str(tests_path / "dp-compressed.pb") + INPUT = str(tests_path / "input.json") + jdata = j_loader(str(tests_path / os.path.join("model_compression", "input.json"))) + jdata["training"]["training_data"]["systems"] = data_file + jdata["training"]["validation_data"]["systems"] = data_file + jdata["model"]["descriptor"]["type_one_side"] = True + jdata["model"]["descriptor"]["exclude_types"] = [[0, 0]] + with open(INPUT, "w") as fp: + json.dump(jdata, fp, indent=4) + + ret = _subprocess_run("dp train " + INPUT) + np.testing.assert_equal(ret, 0, 'DP train failed!') + ret = _subprocess_run("dp freeze -o " + frozen_model) + np.testing.assert_equal(ret, 0, 'DP freeze failed!') + ret = _subprocess_run("dp compress " + " -i " + frozen_model + " -o " + compressed_model) + np.testing.assert_equal(ret, 0, 'DP model compression failed!') + return INPUT, frozen_model, compressed_model + +INPUT, FROZEN_MODEL, COMPRESSED_MODEL = _init_models() + +class TestDeepPotAPBCTypeOneSideExcludeTypes(unittest.TestCase) : + @classmethod + def setUpClass(self): + self.dp_original = DeepPot(FROZEN_MODEL) + self.dp_compressed = DeepPot(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 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 + np.testing.assert_almost_equal(ff0, ff1, default_places) + np.testing.assert_almost_equal(ee0, ee1, default_places) + np.testing.assert_almost_equal(vv0, vv1, 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 + np.testing.assert_almost_equal(ff0, ff1, default_places) + np.testing.assert_almost_equal(ae0, ae1, default_places) + np.testing.assert_almost_equal(av0, av1, default_places) + np.testing.assert_almost_equal(ee0, ee1, default_places) + np.testing.assert_almost_equal(vv0, vv1, 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 + np.testing.assert_almost_equal(ff0, ff1, default_places) + np.testing.assert_almost_equal(ae0, ae1, default_places) + np.testing.assert_almost_equal(av0, av1, default_places) + np.testing.assert_almost_equal(ee0, ee1, default_places) + np.testing.assert_almost_equal(vv0, vv1, default_places) diff --git a/source/tests/test_type_one_side.py b/source/tests/test_type_one_side.py new file mode 100644 index 0000000000..aa1548f637 --- /dev/null +++ b/source/tests/test_type_one_side.py @@ -0,0 +1,135 @@ + +import numpy as np +from deepmd.env import tf +from common import gen_data, j_loader + +from deepmd.utils.data_system import DataSystem +from deepmd.descriptor import Descriptor +from deepmd.common import j_must_have + +GLOBAL_ENER_FLOAT_PRECISION = tf.float64 +GLOBAL_TF_FLOAT_PRECISION = tf.float64 +GLOBAL_NP_FLOAT_PRECISION = np.float64 + +class TestModel(tf.test.TestCase): + def setUp(self): + gen_data(nframes=2) + + def test_descriptor_one_side_exclude_types(self): + """When we enable type_one_side, the descriptor should be the same + for different types, when its environments are the same. + + Here we generates two data. The only difference is the type: + (1) 0 1 1 1 1 1 + (2) 1 1 1 1 1 1 + + When type_one_side is true, the first atom should have the same descriptor. + Otherwise, it should be different (with random initial variables). We test + both situation. + """ + jfile = 'water_se_a.json' + jdata = j_loader(jfile) + + systems = j_must_have(jdata, 'systems') + set_pfx = j_must_have(jdata, 'set_prefix') + batch_size = j_must_have(jdata, 'batch_size') + test_size = j_must_have(jdata, 'numb_test') + batch_size = 1 + test_size = 1 + rcut = j_must_have (jdata['model']['descriptor'], 'rcut') + sel = j_must_have (jdata['model']['descriptor'], 'sel') + ntypes=len(sel) + + data = DataSystem(systems, set_pfx, batch_size, test_size, rcut, run_opt = None) + + test_data = data.get_test () + numb_test = 1 + + # set parameters + jdata['model']['descriptor']['neuron'] = [5, 5, 5] + jdata['model']['descriptor']['axis_neuron'] = 2 + jdata['model']['descriptor']['type_one_side'] = True + jdata['model']['descriptor']['exclude_types'] = [[0, 0]] + + t_prop_c = tf.placeholder(tf.float32, [5], name='t_prop_c') + t_coord = tf.placeholder(GLOBAL_TF_FLOAT_PRECISION, [None], name='i_coord') + t_type = tf.placeholder(tf.int32, [None], name='i_type') + t_natoms = tf.placeholder(tf.int32, [ntypes+2], name='i_natoms') + t_box = tf.placeholder(GLOBAL_TF_FLOAT_PRECISION, [None, 9], name='i_box') + t_mesh = tf.placeholder(tf.int32, [None], name='i_mesh') + is_training = tf.placeholder(tf.bool) + + + # successful + descrpt = Descriptor(**jdata['model']['descriptor']) + dout \ + = descrpt.build( + t_coord, + t_type, + t_natoms, + t_box, + t_mesh, + {}, + reuse = False, + suffix = "_se_a_1side_exclude_types" + ) + # failed + descrpt_failed = Descriptor(**{**jdata['model']['descriptor'], "type_one_side": False}) + dout_failed \ + = descrpt_failed.build( + t_coord, + t_type, + t_natoms, + t_box, + t_mesh, + {}, + reuse = False, + suffix = "_se_a_1side_exclude_types_failed" + ) + + feed_dict_test1 = {t_prop_c: test_data['prop_c'], + t_coord: np.reshape(test_data['coord'] [:numb_test, :], [-1]), + t_box: test_data['box'] [:numb_test, :], + t_type: np.reshape(test_data['type'] [:numb_test, :], [-1]), + t_natoms: test_data['natoms_vec'], + t_mesh: test_data['default_mesh'], + is_training: False} + feed_dict_test2 = feed_dict_test1.copy() + # original type: 0 0 1 1 1 1 + # current: 0 1 1 1 1 1 + # current: 1 1 1 1 1 1 + new_natoms1 = test_data['natoms_vec'].copy() + new_natoms1[2] = 1 + new_natoms1[3] = 5 + new_type1 = test_data['type'].copy() + new_type1[:numb_test, 0] = 0 + new_type1[:numb_test, 1:6] = 1 + new_natoms2 = test_data['natoms_vec'].copy() + new_natoms2[2] = 0 + new_natoms2[3] = 6 + new_type2 = test_data['type'].copy() + new_type2[:numb_test] = 1 + feed_dict_test1[t_type] = np.reshape(new_type1[:numb_test, :], [-1]) + feed_dict_test1[t_natoms] = new_natoms1 + feed_dict_test2[t_type] = np.reshape(new_type2[:numb_test, :], [-1]) + feed_dict_test2[t_natoms] = new_natoms2 + print(feed_dict_test1,feed_dict_test2) + + with self.test_session() as sess: + sess.run(tf.global_variables_initializer()) + [model_dout1] = sess.run([dout], + feed_dict = feed_dict_test1) + [model_dout2] = sess.run([dout], + feed_dict = feed_dict_test2) + [model_dout1_failed] = sess.run([dout_failed], + feed_dict = feed_dict_test1) + [model_dout2_failed] = sess.run([dout_failed], + feed_dict = feed_dict_test2) + model_dout1 = model_dout1.reshape([6, -1]) + model_dout2 = model_dout2.reshape([6, -1]) + model_dout1_failed = model_dout1_failed.reshape([6, -1]) + model_dout2_failed = model_dout2_failed.reshape([6, -1]) + + np.testing.assert_almost_equal(model_dout1[0], model_dout2[0], 10) + with self.assertRaises(AssertionError): + np.testing.assert_almost_equal(model_dout1_failed[0], model_dout2_failed[0], 10) From 3d68634ee672415274c438cc14a61e16538579fe Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Sun, 13 Mar 2022 20:37:58 -0400 Subject: [PATCH 7/7] rename model name to prevent conflict --- ...test_model_compression_se_a_type_one_side_exclude_types.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/source/tests/test_model_compression_se_a_type_one_side_exclude_types.py b/source/tests/test_model_compression_se_a_type_one_side_exclude_types.py index d29fc6b1fe..b223742d2b 100644 --- a/source/tests/test_model_compression_se_a_type_one_side_exclude_types.py +++ b/source/tests/test_model_compression_se_a_type_one_side_exclude_types.py @@ -32,8 +32,8 @@ def _subprocess_run(command): def _init_models(): data_file = str(tests_path / os.path.join("model_compression", "data")) - frozen_model = str(tests_path / "dp-original.pb") - compressed_model = str(tests_path / "dp-compressed.pb") + frozen_model = str(tests_path / "dp-original-type-one-side-exclude-types.pb") + compressed_model = str(tests_path / "dp-compressed-type-one-side-exclude-types.pb") INPUT = str(tests_path / "input.json") jdata = j_loader(str(tests_path / os.path.join("model_compression", "input.json"))) jdata["training"]["training_data"]["systems"] = data_file