From d5679950388e7a0ebd912c0cd77f44c4b167e40c Mon Sep 17 00:00:00 2001 From: Han Wang Date: Mon, 9 Oct 2023 17:01:20 +0800 Subject: [PATCH 1/3] fix the smoothness issue of se_attn_v2 --- deepmd/descriptor/se_atten.py | 23 ++++- source/tests/common.py | 72 ++++++++++++++ source/tests/test_model_se_atten.py | 139 ++++++++++++++++++++++++++++ 3 files changed, 230 insertions(+), 4 deletions(-) diff --git a/deepmd/descriptor/se_atten.py b/deepmd/descriptor/se_atten.py index 90cc86c4d9..102bc5ec6f 100644 --- a/deepmd/descriptor/se_atten.py +++ b/deepmd/descriptor/se_atten.py @@ -564,6 +564,8 @@ def build( self.filter_precision, ) self.negative_mask = -(2 << 32) * (1.0 - self.nmask) + # hard coding the magnitude of attention weight shift + self.smth_attn_w_shift = 20.0 # only used when tensorboard was set as true tf.summary.histogram("descrpt", self.descrpt) tf.summary.histogram("rij", self.rij) @@ -599,7 +601,7 @@ def build( ) self.recovered_r = ( tf.reshape( - tf.slice(tf.reshape(self.descrpt, [-1, 4]), [0, 0], [-1, 1]), + tf.slice(tf.reshape(self.descrpt_reshape, [-1, 4]), [0, 0], [-1, 1]), [-1, natoms[0], self.sel_all_a[0]], ) * self.std_looked_up @@ -865,10 +867,23 @@ def _scaled_dot_attn( save_weights=True, ): attn = tf.matmul(Q / temperature, K, transpose_b=True) - attn *= self.nmask - attn += self.negative_mask + if self.smooth: + # (nb x nloc) x nsel + nsel = self.sel_all_a[0] + attn = ((attn + self.smth_attn_w_shift) * + tf.reshape(self.recovered_switch, [-1,1,nsel]) * + tf.reshape(self.recovered_switch, [-1,nsel,1]) - + self.smth_attn_w_shift) + else: + attn *= self.nmask + attn += self.negative_mask attn = tf.nn.softmax(attn, axis=-1) - attn *= tf.reshape(self.nmask, [-1, attn.shape[-1], 1]) + if self.smooth: + attn = (attn * + tf.reshape(self.recovered_switch, [-1,1,nsel]) * + tf.reshape(self.recovered_switch, [-1,nsel,1])) + else: + attn *= tf.reshape(self.nmask, [-1, attn.shape[-1], 1]) if save_weights: self.attn_weight[layer] = attn[0] # atom 0 if dotr: diff --git a/source/tests/common.py b/source/tests/common.py index f8ed23df03..40724dc406 100644 --- a/source/tests/common.py +++ b/source/tests/common.py @@ -530,6 +530,78 @@ def strerch_box(old_coord, old_box, new_box): return ncoord.reshape(old_coord.shape) +def finite_difference_fv(sess, energy, feed_dict, t_coord, t_box, delta=1e-6): + """for energy models, compute f, v by finite difference + """ + base_dict = feed_dict.copy() + coord0 = base_dict.pop(t_coord) + box0 = base_dict.pop(t_box) + fdf = -finite_difference( + lambda coord: sess.run( + energy, feed_dict={**base_dict, t_coord: coord, t_box: box0} + ).reshape(-1), + coord0, + delta=delta, + ).reshape(-1) + fdv = -( + finite_difference( + lambda box: sess.run( + energy, + feed_dict={ + **base_dict, + t_coord: strerch_box(coord0, box0, box), + t_box: box, + }, + ).reshape(-1), + box0, + delta=delta, + ) + .reshape([-1, 3, 3]) + .transpose(0, 2, 1) + @ box0.reshape(3, 3) + ).reshape(-1) + return fdf, fdv + + +def check_continuity(f, cc, rcut, delta): + """coord[0:2] to [[0, 0, 0], [rcut+-.5*delta, 0, 0]] + """ + cc = cc.reshape([-1,3]) + cc0 = np.copy(cc) + cc1 = np.copy(cc) + cc0[:2,:] = np.array([ + 0.0, 0.0, 0.0, + rcut-0.5*delta, 0.0, 0.0, + ]).reshape([-1,3]) + cc1[:2,:] = np.array([ + 0.0, 0.0, 0.0, + rcut+0.5*delta, 0.0, 0.0, + ]).reshape([-1,3]) + return f(cc0.reshape(-1)), f(cc1.reshape(-1)) + + +def check_smooth_efv(sess, energy, force, virial, feed_dict, t_coord, rcut, delta=1e-5): + """check the smoothness of e, f and v + the returned values are de, df, dv + de[0] are supposed to be closed to de[1] + df[0] are supposed to be closed to df[1] + dv[0] are supposed to be closed to dv[1] + """ + base_dict = feed_dict.copy() + coord0 = base_dict.pop(t_coord) + [fe, ff, fv] = [ + lambda coord: sess.run( + ii, feed_dict={**base_dict, t_coord: coord} + ).reshape(-1) + for ii in [energy, force, virial] + ] + [de, df, dv] = [ + check_continuity(ii, coord0, rcut, delta=delta) + for ii in [fe, ff, fv] + ] + return de, df, dv + + def run_dp(cmd: str) -> int: """Run DP directly from the entry point instead of the subprocess. diff --git a/source/tests/test_model_se_atten.py b/source/tests/test_model_se_atten.py index 445959ceb2..e674c88998 100644 --- a/source/tests/test_model_se_atten.py +++ b/source/tests/test_model_se_atten.py @@ -28,6 +28,7 @@ from deepmd.utils.type_embed import ( TypeEmbedNet, ) +from common import finite_difference_fv, check_smooth_efv GLOBAL_ENER_FLOAT_PRECISION = tf.float64 GLOBAL_TF_FLOAT_PRECISION = tf.float64 @@ -726,3 +727,141 @@ def test_stripped_type_embedding_exclude_types(self): np.testing.assert_almost_equal(des[:, 0:2], 0.0, 10) with self.assertRaises(AssertionError): np.testing.assert_almost_equal(des[:, 2:6], 0.0, 10) + + + def test_smoothness_of_stripped_type_embedding_smooth_model(self): + """test: auto-diff, continuity of e,f,v + """ + jfile = "water_se_atten.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 + stop_batch = j_must_have(jdata, "stop_batch") + rcut = j_must_have(jdata["model"]["descriptor"], "rcut") + + data = DataSystem(systems, set_pfx, batch_size, test_size, rcut, run_opt=None) + + test_data = data.get_test() + numb_test = 1 + + jdata["model"]["descriptor"].pop("type", None) + jdata["model"]["descriptor"]["ntypes"] = 2 + jdata["model"]["descriptor"]["stripped_type_embedding"] = True + jdata["model"]["descriptor"]["smooth_type_embdding"] = True + jdata["model"]["descriptor"]["attn_layer"] = 1 + descrpt = DescrptSeAtten(**jdata["model"]["descriptor"], uniform_seed=True) + jdata["model"]["fitting_net"]["descrpt"] = descrpt + fitting = EnerFitting(**jdata["model"]["fitting_net"], uniform_seed=True) + typeebd_param = jdata["model"]["type_embedding"] + typeebd = TypeEmbedNet( + neuron=typeebd_param["neuron"], + activation_function=None, + resnet_dt=typeebd_param["resnet_dt"], + seed=typeebd_param["seed"], + uniform_seed=True, + padding=True, + ) + model = EnerModel(descrpt, fitting, typeebd) + + # model._compute_dstats([test_data['coord']], [test_data['box']], [test_data['type']], [test_data['natoms_vec']], [test_data['default_mesh']]) + input_data = { + "coord": [test_data["coord"]], + "box": [test_data["box"]], + "type": [test_data["type"]], + "natoms_vec": [test_data["natoms_vec"]], + "default_mesh": [test_data["default_mesh"]], + } + model._compute_input_stat(input_data) + model.descrpt.bias_atom_e = data.compute_energy_shift() + + t_prop_c = tf.placeholder(tf.float32, [5], name="t_prop_c") + t_energy = tf.placeholder(GLOBAL_ENER_FLOAT_PRECISION, [None], name="t_energy") + t_force = tf.placeholder(GLOBAL_TF_FLOAT_PRECISION, [None], name="t_force") + t_virial = tf.placeholder(GLOBAL_TF_FLOAT_PRECISION, [None], name="t_virial") + t_atom_ener = tf.placeholder( + GLOBAL_TF_FLOAT_PRECISION, [None], name="t_atom_ener" + ) + 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, [model.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) + t_fparam = None + inputs_dict = {} + + model_pred = model.build( + t_coord, + t_type, + t_natoms, + t_box, + t_mesh, + inputs_dict, + suffix=self.filename + + "-" + + inspect.stack()[0][3] + + "test_model_se_atten_model_compressible", + reuse=False, + ) + energy = model_pred["energy"] + force = model_pred["force"] + virial = model_pred["virial"] + atom_ener = model_pred["atom_ener"] + + feed_dict_test = { + t_prop_c: test_data["prop_c"], + t_energy: test_data["energy"][:numb_test], + t_force: np.reshape(test_data["force"][:numb_test, :], [-1]), + t_virial: np.reshape(test_data["virial"][:numb_test, :], [-1]), + t_atom_ener: np.reshape(test_data["atom_ener"][:numb_test, :], [-1]), + 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, + } + sess = self.cached_session().__enter__() + sess.run(tf.global_variables_initializer()) + [pe, pf, pv] = sess.run([energy, force, virial], feed_dict=feed_dict_test) + pf, pv = pf.reshape(-1), pv.reshape(-1) + + eps = 1e-4 + delta = 1e-5 + fdf, fdv = finite_difference_fv( + sess, energy, feed_dict_test, t_coord, t_box, delta=eps) + np.testing.assert_allclose(pf, fdf, delta) + np.testing.assert_allclose(pv, fdv, delta) + + tested_eps = [1e-3, 1e-4, 1e-5, 1e-6, 1e-7] + for eps in tested_eps: + deltae = eps + deltad = eps + de, df, dv = check_smooth_efv( + sess, energy, force, virial, + feed_dict_test, t_coord, + jdata["model"]["descriptor"]["rcut"], + delta=eps, + ) + np.testing.assert_allclose(de[0], de[1], rtol=0, atol=deltae) + np.testing.assert_allclose(df[0], df[1], rtol=0, atol=deltad) + np.testing.assert_allclose(dv[0], dv[1], rtol=0, atol=deltad) + + for eps in tested_eps: + deltae = 5.*eps + deltad = 5.*eps + de, df, dv = check_smooth_efv( + sess, energy, force, virial, + feed_dict_test, t_coord, + jdata["model"]["descriptor"]["rcut_smth"], + delta=eps, + ) + np.testing.assert_allclose(de[0], de[1], rtol=0, atol=deltae) + np.testing.assert_allclose(df[0], df[1], rtol=0, atol=deltad) + np.testing.assert_allclose(dv[0], dv[1], rtol=0, atol=deltad) + From 7ceabd2a693fa16536ea42d209439129ac161702 Mon Sep 17 00:00:00 2001 From: Han Wang Date: Fri, 13 Oct 2023 11:46:18 +0800 Subject: [PATCH 2/3] add uts for the compression of smooth se_atten --- .../tests/test_model_compression_se_atten.py | 18 +++++++++++------- source/tests/test_model_se_atten.py | 2 ++ 2 files changed, 13 insertions(+), 7 deletions(-) diff --git a/source/tests/test_model_compression_se_atten.py b/source/tests/test_model_compression_se_atten.py index 30d78669ce..ea28f5c56c 100644 --- a/source/tests/test_model_compression_se_atten.py +++ b/source/tests/test_model_compression_se_atten.py @@ -46,10 +46,11 @@ def _subprocess_run(command): # - type embedding FP32, se_atten FP64 # - type embedding FP32, se_atten FP32 tests = [ - {"se_atten precision": "float64", "type embedding precision": "float64"}, - {"se_atten precision": "float64", "type embedding precision": "float32"}, - {"se_atten precision": "float32", "type embedding precision": "float64"}, - {"se_atten precision": "float32", "type embedding precision": "float32"}, + {"se_atten precision": "float64", "type embedding precision": "float64", "smooth_type_embdding": True}, + {"se_atten precision": "float64", "type embedding precision": "float64", "smooth_type_embdding": False}, + {"se_atten precision": "float64", "type embedding precision": "float32", "smooth_type_embdding": True}, + {"se_atten precision": "float32", "type embedding precision": "float64", "smooth_type_embdding": True}, + {"se_atten precision": "float32", "type embedding precision": "float32", "smooth_type_embdding": True}, ] @@ -73,6 +74,9 @@ def _init_models(): jdata["model"]["descriptor"]["stripped_type_embedding"] = True jdata["model"]["descriptor"]["sel"] = 120 jdata["model"]["descriptor"]["attn_layer"] = 0 + jdata["model"]["descriptor"]["smooth_type_embdding"] = tests[i][ + "smooth_type_embdding" + ] jdata["model"]["type_embedding"] = {} jdata["model"]["type_embedding"]["precision"] = tests[i][ "type embedding precision" @@ -479,9 +483,9 @@ def test_1frame(self): 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) + np.testing.assert_almost_equal(ff0, ff1, default_places, err_msg=str(tests[i])) + np.testing.assert_almost_equal(ee0, ee1, default_places, err_msg=str(tests[i])) + np.testing.assert_almost_equal(vv0, vv1, default_places, err_msg=str(tests[i])) def test_1frame_atm(self): for i in range(len(tests)): diff --git a/source/tests/test_model_se_atten.py b/source/tests/test_model_se_atten.py index e674c88998..605bf55404 100644 --- a/source/tests/test_model_se_atten.py +++ b/source/tests/test_model_se_atten.py @@ -754,6 +754,8 @@ def test_smoothness_of_stripped_type_embedding_smooth_model(self): jdata["model"]["descriptor"]["stripped_type_embedding"] = True jdata["model"]["descriptor"]["smooth_type_embdding"] = True jdata["model"]["descriptor"]["attn_layer"] = 1 + jdata["model"]["descriptor"]["rcut"] = 6.0 + jdata["model"]["descriptor"]["rcut_smth"] = 4.0 descrpt = DescrptSeAtten(**jdata["model"]["descriptor"], uniform_seed=True) jdata["model"]["fitting_net"]["descrpt"] = descrpt fitting = EnerFitting(**jdata["model"]["fitting_net"], uniform_seed=True) From 175c2135152e148b2b14c9c96c75fdbe05a0fd29 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 13 Oct 2023 03:50:33 +0000 Subject: [PATCH 3/3] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- source/tests/common.py | 1 + .../tests/test_model_compression_se_atten.py | 44 +++++++++++++++---- source/tests/test_model_se_atten.py | 1 - 3 files changed, 36 insertions(+), 10 deletions(-) diff --git a/source/tests/common.py b/source/tests/common.py index 67305ed10c..9af324896f 100644 --- a/source/tests/common.py +++ b/source/tests/common.py @@ -608,6 +608,7 @@ def check_smooth_efv(sess, energy, force, virial, feed_dict, t_coord, rcut, delt ] return de, df, dv + def run_dp(cmd: str) -> int: """Run DP directly from the entry point instead of the subprocess. diff --git a/source/tests/test_model_compression_se_atten.py b/source/tests/test_model_compression_se_atten.py index ea28f5c56c..6bab1a3881 100644 --- a/source/tests/test_model_compression_se_atten.py +++ b/source/tests/test_model_compression_se_atten.py @@ -46,11 +46,31 @@ def _subprocess_run(command): # - type embedding FP32, se_atten FP64 # - type embedding FP32, se_atten FP32 tests = [ - {"se_atten precision": "float64", "type embedding precision": "float64", "smooth_type_embdding": True}, - {"se_atten precision": "float64", "type embedding precision": "float64", "smooth_type_embdding": False}, - {"se_atten precision": "float64", "type embedding precision": "float32", "smooth_type_embdding": True}, - {"se_atten precision": "float32", "type embedding precision": "float64", "smooth_type_embdding": True}, - {"se_atten precision": "float32", "type embedding precision": "float32", "smooth_type_embdding": True}, + { + "se_atten precision": "float64", + "type embedding precision": "float64", + "smooth_type_embdding": True, + }, + { + "se_atten precision": "float64", + "type embedding precision": "float64", + "smooth_type_embdding": False, + }, + { + "se_atten precision": "float64", + "type embedding precision": "float32", + "smooth_type_embdding": True, + }, + { + "se_atten precision": "float32", + "type embedding precision": "float64", + "smooth_type_embdding": True, + }, + { + "se_atten precision": "float32", + "type embedding precision": "float32", + "smooth_type_embdding": True, + }, ] @@ -75,7 +95,7 @@ def _init_models(): jdata["model"]["descriptor"]["sel"] = 120 jdata["model"]["descriptor"]["attn_layer"] = 0 jdata["model"]["descriptor"]["smooth_type_embdding"] = tests[i][ - "smooth_type_embdding" + "smooth_type_embdding" ] jdata["model"]["type_embedding"] = {} jdata["model"]["type_embedding"]["precision"] = tests[i][ @@ -483,9 +503,15 @@ def test_1frame(self): self.assertEqual(ff1.shape, (nframes, natoms, 3)) self.assertEqual(vv1.shape, (nframes, 9)) # check values - np.testing.assert_almost_equal(ff0, ff1, default_places, err_msg=str(tests[i])) - np.testing.assert_almost_equal(ee0, ee1, default_places, err_msg=str(tests[i])) - np.testing.assert_almost_equal(vv0, vv1, default_places, err_msg=str(tests[i])) + np.testing.assert_almost_equal( + ff0, ff1, default_places, err_msg=str(tests[i]) + ) + np.testing.assert_almost_equal( + ee0, ee1, default_places, err_msg=str(tests[i]) + ) + np.testing.assert_almost_equal( + vv0, vv1, default_places, err_msg=str(tests[i]) + ) def test_1frame_atm(self): for i in range(len(tests)): diff --git a/source/tests/test_model_se_atten.py b/source/tests/test_model_se_atten.py index c8b5b5b3f1..5417201a9f 100644 --- a/source/tests/test_model_se_atten.py +++ b/source/tests/test_model_se_atten.py @@ -30,7 +30,6 @@ from deepmd.utils.type_embed import ( TypeEmbedNet, ) -from common import finite_difference_fv, check_smooth_efv GLOBAL_ENER_FLOAT_PRECISION = tf.float64 GLOBAL_TF_FLOAT_PRECISION = tf.float64