From bda1f7bf67875e0c03ef79271acae450e51d4de4 Mon Sep 17 00:00:00 2001 From: Duo <50307526+iProzd@users.noreply.github.com> Date: Fri, 31 May 2024 17:30:56 +0800 Subject: [PATCH 01/10] fix: correct `exclude_types` in descriptors --- deepmd/dpmodel/descriptor/dpa1.py | 6 +- deepmd/dpmodel/descriptor/repformers.py | 2 +- deepmd/dpmodel/descriptor/se_e2_a.py | 5 + deepmd/dpmodel/descriptor/se_r.py | 5 + deepmd/dpmodel/descriptor/se_t.py | 5 + deepmd/pt/model/descriptor/repformers.py | 2 +- deepmd/pt/model/descriptor/se_atten.py | 13 +- deepmd/tf/descriptor/se_atten.py | 19 ++- source/tests/pt/model/test_exclude_types.py | 151 ++++++++++++++++++ source/tests/universal/common/cases/cases.py | 132 +++++++++++++++ .../common/cases/descriptor/__init__.py | 1 + .../common/cases/descriptor/descriptor.py | 11 ++ .../common/cases/descriptor/utils.py | 95 +++++++++++ .../universal/dpmodel/descriptor/__init__.py | 1 + .../dpmodel/descriptor/test_descriptor.py | 84 ++++++++++ .../tests/universal/pt/descriptor/__init__.py | 1 + .../pt/descriptor/test_descriptor.py | 84 ++++++++++ 17 files changed, 605 insertions(+), 12 deletions(-) create mode 100644 source/tests/pt/model/test_exclude_types.py create mode 100644 source/tests/universal/common/cases/cases.py create mode 100644 source/tests/universal/common/cases/descriptor/__init__.py create mode 100644 source/tests/universal/common/cases/descriptor/descriptor.py create mode 100644 source/tests/universal/common/cases/descriptor/utils.py create mode 100644 source/tests/universal/dpmodel/descriptor/__init__.py create mode 100644 source/tests/universal/dpmodel/descriptor/test_descriptor.py create mode 100644 source/tests/universal/pt/descriptor/__init__.py create mode 100644 source/tests/universal/pt/descriptor/test_descriptor.py diff --git a/deepmd/dpmodel/descriptor/dpa1.py b/deepmd/dpmodel/descriptor/dpa1.py index 42929ca1d3..d9f8570cd0 100644 --- a/deepmd/dpmodel/descriptor/dpa1.py +++ b/deepmd/dpmodel/descriptor/dpa1.py @@ -810,6 +810,8 @@ def call( nf, nloc, nnei, _ = dmatrix.shape exclude_mask = self.emask.build_type_exclude_mask(nlist, atype_ext) # nfnl x nnei + exclude_mask = exclude_mask.reshape(nf * nloc, nnei) + # nfnl x nnei nlist = nlist.reshape(nf * nloc, nnei) # nfnl x nnei x 4 dmatrix = dmatrix.reshape(nf * nloc, nnei, 4) @@ -821,6 +823,8 @@ def call( atype_embd_nnei = np.tile(atype_embd[:, np.newaxis, :], (1, nnei, 1)) # nfnl x nnei nlist_mask = nlist != -1 + # nfnl x nnei + nlist_mask = nlist_mask * exclude_mask.astype(bool) # nfnl x nnei x 1 sw = np.where(nlist_mask[:, :, None], sw, 0.0) nlist_masked = np.where(nlist_mask, nlist, 0) @@ -830,8 +834,6 @@ def call( nf * nloc, nnei, self.tebd_dim ) ng = self.neuron[-1] - # nfnl x nnei - exclude_mask = exclude_mask.reshape(nf * nloc, nnei) # nfnl x nnei x 4 rr = dmatrix.reshape(nf * nloc, nnei, 4) rr = rr * exclude_mask[:, :, None] diff --git a/deepmd/dpmodel/descriptor/repformers.py b/deepmd/dpmodel/descriptor/repformers.py index f0275b44b1..8e18913296 100644 --- a/deepmd/dpmodel/descriptor/repformers.py +++ b/deepmd/dpmodel/descriptor/repformers.py @@ -346,7 +346,7 @@ def call( mapping: Optional[np.ndarray] = None, ): exclude_mask = self.emask.build_type_exclude_mask(nlist, atype_ext) - nlist = nlist * exclude_mask + nlist = np.where(exclude_mask, nlist, -1) # nf x nloc x nnei x 4 dmatrix, diff, sw = self.env_mat.call( coord_ext, atype_ext, nlist, self.mean, self.stddev diff --git a/deepmd/dpmodel/descriptor/se_e2_a.py b/deepmd/dpmodel/descriptor/se_e2_a.py index 2b6c3a843e..ed332615a6 100644 --- a/deepmd/dpmodel/descriptor/se_e2_a.py +++ b/deepmd/dpmodel/descriptor/se_e2_a.py @@ -124,6 +124,9 @@ class DescrptSeA(NativeOP, BaseDescriptor): The precision of the embedding net parameters. Supported options are |PRECISION| spin The deepspin object. + ntypes : int + Number of element types. + Not used in this descriptor, only to be compat with input. Limitations ----------- @@ -157,9 +160,11 @@ def __init__( activation_function: str = "tanh", precision: str = DEFAULT_PRECISION, spin: Optional[Any] = None, + ntypes: Optional[int] = None, # to be compat with input # consistent with argcheck, not used though seed: Optional[int] = None, ) -> None: + del ntypes ## seed, uniform_seed, not included. if spin is not None: raise NotImplementedError("spin is not implemented") diff --git a/deepmd/dpmodel/descriptor/se_r.py b/deepmd/dpmodel/descriptor/se_r.py index e4a0a80657..134d413de6 100644 --- a/deepmd/dpmodel/descriptor/se_r.py +++ b/deepmd/dpmodel/descriptor/se_r.py @@ -81,6 +81,9 @@ class DescrptSeR(NativeOP, BaseDescriptor): The precision of the embedding net parameters. Supported options are |PRECISION| spin The deepspin object. + ntypes : int + Number of element types. + Not used in this descriptor, only to be compat with input. Limitations ----------- @@ -113,9 +116,11 @@ def __init__( activation_function: str = "tanh", precision: str = DEFAULT_PRECISION, spin: Optional[Any] = None, + ntypes: Optional[int] = None, # to be compat with input # consistent with argcheck, not used though seed: Optional[int] = None, ) -> None: + del ntypes ## seed, uniform_seed, not included. if not type_one_side: raise NotImplementedError("type_one_side == False not implemented") diff --git a/deepmd/dpmodel/descriptor/se_t.py b/deepmd/dpmodel/descriptor/se_t.py index eac6a9640e..9f953b87d5 100644 --- a/deepmd/dpmodel/descriptor/se_t.py +++ b/deepmd/dpmodel/descriptor/se_t.py @@ -85,6 +85,9 @@ class DescrptSeT(NativeOP, BaseDescriptor): If the weights of embedding net are trainable. seed : int, Optional Random seed for initializing the network parameters. + ntypes : int + Number of element types. + Not used in this descriptor, only to be compat with input. """ def __init__( @@ -101,7 +104,9 @@ def __init__( precision: str = DEFAULT_PRECISION, trainable: bool = True, seed: Optional[int] = None, + ntypes: Optional[int] = None, # to be compat with input ) -> None: + del ntypes self.rcut = rcut self.rcut_smth = rcut_smth self.sel = sel diff --git a/deepmd/pt/model/descriptor/repformers.py b/deepmd/pt/model/descriptor/repformers.py index a66693653e..2e091d6bb8 100644 --- a/deepmd/pt/model/descriptor/repformers.py +++ b/deepmd/pt/model/descriptor/repformers.py @@ -410,7 +410,7 @@ def forward( atype = extended_atype[:, :nloc] # nb x nloc x nnei exclude_mask = self.emask(nlist, extended_atype) - nlist = nlist * exclude_mask + nlist = torch.where(exclude_mask != 0, nlist, -1) # nb x nloc x nnei x 4, nb x nloc x nnei x 3, nb x nloc x nnei x 1 dmatrix, diff, sw = prod_env_mat( extended_coord, diff --git a/deepmd/pt/model/descriptor/se_atten.py b/deepmd/pt/model/descriptor/se_atten.py index a59eaca409..d26578dc74 100644 --- a/deepmd/pt/model/descriptor/se_atten.py +++ b/deepmd/pt/model/descriptor/se_atten.py @@ -481,8 +481,6 @@ def forward( nlist_mask = nlist != -1 nlist = torch.where(nlist == -1, 0, nlist) sw = torch.squeeze(sw, -1) - # beyond the cutoff sw should be 0.0 - sw = sw.masked_fill(~nlist_mask, 0.0) # nf x nloc x nt -> nf x nloc x nnei x nt atype_tebd = extended_atype_embd[:, :nloc, :] atype_tebd_nnei = atype_tebd.unsqueeze(2).expand(-1, -1, self.nnei, -1) @@ -495,8 +493,17 @@ def forward( atype_tebd_nlist = torch.gather(atype_tebd_ext, dim=1, index=index) # nb x nloc x nnei x nt atype_tebd_nlist = atype_tebd_nlist.view(nb, nloc, nnei, nt) + # nb x nloc x nnei + exclude_mask = self.emask(nlist, extended_atype) + nlist_mask = torch.where( + exclude_mask != 0, + nlist_mask, + torch.zeros_like(nlist_mask, dtype=torch.bool), + ) + # beyond the cutoff sw should be 0.0 + sw = sw.masked_fill(~nlist_mask, 0.0) # (nb x nloc) x nnei - exclude_mask = self.emask(nlist, extended_atype).view(nb * nloc, nnei) + exclude_mask = exclude_mask.view(nb * nloc, nnei) if self.old_impl: assert self.filter_layers_old is not None dmatrix = dmatrix.view( diff --git a/deepmd/tf/descriptor/se_atten.py b/deepmd/tf/descriptor/se_atten.py index 6d3cfeaa6e..c05d31ba60 100644 --- a/deepmd/tf/descriptor/se_atten.py +++ b/deepmd/tf/descriptor/se_atten.py @@ -718,6 +718,18 @@ def _pass_filter( tf.shape(inputs_i)[0], self.nei_type_vec, # extra input for atten ) + # (nsamples * natoms * nnei, 1) + nei_exclude_mask = tf.slice( + tf.reshape(tf.cast(mask, self.filter_precision), [-1, 4]), + [0, 0], + [-1, 1], + ) + # (nsamples * natoms, 1, nnei) + self.nmask *= tf.reshape( + nei_exclude_mask, + [-1, 1, self.sel_all_a[0]], + ) + self.negative_mask = -(2 << 32) * (1.0 - self.nmask) if self.smooth: inputs_i = tf.where( tf.cast(mask, tf.bool), @@ -727,12 +739,9 @@ def _pass_filter( tf.reshape(self.avg_looked_up, [-1, 1]), [1, self.ndescrpt] ), ) + # (nsamples, natoms, nnei) self.recovered_switch *= tf.reshape( - tf.slice( - tf.reshape(tf.cast(mask, self.filter_precision), [-1, 4]), - [0, 0], - [-1, 1], - ), + nei_exclude_mask, [-1, natoms[0], self.sel_all_a[0]], ) else: diff --git a/source/tests/pt/model/test_exclude_types.py b/source/tests/pt/model/test_exclude_types.py new file mode 100644 index 0000000000..cf7eb55705 --- /dev/null +++ b/source/tests/pt/model/test_exclude_types.py @@ -0,0 +1,151 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import unittest + +import numpy as np +import torch + +from deepmd.pt.model.descriptor import ( + DescrptDPA1, + DescrptDPA2, + DescrptHybrid, + DescrptSeA, + DescrptSeR, + DescrptSeT, +) +from deepmd.pt.utils import ( + PairExcludeMask, + env, +) +from deepmd.pt.utils.env import ( + PRECISION_DICT, +) + +from .test_env_mat import ( + TestCaseSingleFrameWithNlist, +) +from .test_mlp import ( + get_tols, +) + +dtype = env.GLOBAL_PT_FLOAT_PRECISION + + +class DescrptExcludeTypes(TestCaseSingleFrameWithNlist): + def setUp(self): + TestCaseSingleFrameWithNlist.setUp(self) + self.input_dict = { + "ntypes": self.nt, + "rcut": self.rcut, + "rcut_smth": self.rcut_smth, + "sel": self.sel, + } + + def test_exclude_types( + self, + ): + dtype = PRECISION_DICT["float64"] + rtol, atol = get_tols("float64") + coord_ext_device = torch.tensor(self.coord_ext, dtype=dtype, device=env.DEVICE) + atype_ext_device = torch.tensor(self.atype_ext, dtype=int, device=env.DEVICE) + nlist_device = torch.tensor(self.nlist, dtype=int, device=env.DEVICE) + mapping_device = torch.tensor(self.mapping, dtype=int, device=env.DEVICE) + dd = self.descrpt(**self.input_dict).to(env.DEVICE) + + for em in [[[0, 1]], [[1, 1]]]: + dd0 = self.descrpt(**self.input_dict, exclude_types=em).to(env.DEVICE) + # only equal when set_davg_zero is True + # dd0.sea.mean = torch.tensor(davg, dtype=dtype, device=env.DEVICE) + # dd0.sea.dstd = torch.tensor(dstd, dtype=dtype, device=env.DEVICE) + dd.load_state_dict(dd0.state_dict()) + + ex_pair = PairExcludeMask(self.nt, em) + pair_mask = ex_pair(nlist_device, atype_ext_device) + # exclude neighbors in the nlist + nlist_exclude = torch.where(pair_mask == 1, nlist_device, -1) + + rd0, _, _, _, sw0 = dd0( + coord_ext_device, + atype_ext_device, + nlist_device, + mapping=mapping_device, + ) + + rd_ex, _, _, _, sw_ex = dd( + coord_ext_device, + atype_ext_device, + nlist_exclude, + mapping=mapping_device, + ) + + np.testing.assert_allclose( + rd0.detach().cpu().numpy(), + rd_ex.detach().cpu().numpy(), + rtol=rtol, + atol=atol, + ) + + +class TestDescrptExcludeTypesSeA(DescrptExcludeTypes, unittest.TestCase): + def setUp(self): + super().setUp() + self.descrpt = DescrptSeA + + +class TestDescrptExcludeTypesSeR(DescrptExcludeTypes, unittest.TestCase): + def setUp(self): + super().setUp() + self.descrpt = DescrptSeR + + +class TestDescrptExcludeTypesSeT(DescrptExcludeTypes, unittest.TestCase): + def setUp(self): + super().setUp() + self.descrpt = DescrptSeT + + +class TestDescrptExcludeTypesDPA1(DescrptExcludeTypes, unittest.TestCase): + def setUp(self): + super().setUp() + self.descrpt = DescrptDPA1 + + +class TestDescrptExcludeTypesDPA2(DescrptExcludeTypes, unittest.TestCase): + def setUp(self): + super().setUp() + self.descrpt = DescrptDPA2 + self.input_dict = { + "ntypes": self.nt, + "repinit": { + "rcut": self.rcut, + "rcut_smth": self.rcut_smth, + "nsel": self.sel_mix, + }, + "repformer": { + "rcut": self.rcut / 2, + "rcut_smth": self.rcut_smth, + "nsel": self.sel_mix[0] // 2, + }, + } + + +class TestDescrptExcludeTypesHybrid(DescrptExcludeTypes, unittest.TestCase): + def setUp(self): + super().setUp() + self.descrpt = DescrptHybrid + ddsub0 = { + "type": "se_e2_a", + "ntypes": self.nt, + "rcut": self.rcut, + "rcut_smth": self.rcut_smth, + "sel": self.sel, + } + ddsub1 = { + "type": "dpa1", + "ntypes": self.nt, + "rcut": self.rcut, + "rcut_smth": self.rcut_smth, + "sel": self.sel_mix, + } + self.input_dict = { + "list": [ddsub0, ddsub1], + } diff --git a/source/tests/universal/common/cases/cases.py b/source/tests/universal/common/cases/cases.py new file mode 100644 index 0000000000..797ab188a8 --- /dev/null +++ b/source/tests/universal/common/cases/cases.py @@ -0,0 +1,132 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import numpy as np + + +class TestCaseSingleFrameWithNlist: + def setUp(self): + # nloc == 3, nall == 4 + self.nloc = 3 + self.nall = 4 + self.nf, self.nt = 2, 2 + self.coord_ext = np.array( + [ + [0, 0, 0], + [0, 1, 0], + [0, 0, 1], + [0, -2, 0], + ], + dtype=np.float64, + ).reshape([1, self.nall, 3]) + self.atype_ext = np.array([0, 0, 1, 0], dtype=int).reshape([1, self.nall]) + self.mapping = np.array([0, 1, 2, 0], dtype=int).reshape([1, self.nall]) + # sel = [5, 2] + self.sel = [5, 2] + self.sel_mix = [7] + self.natoms = [3, 3, 2, 1] + self.nlist = np.array( + [ + [1, 3, -1, -1, -1, 2, -1], + [0, -1, -1, -1, -1, 2, -1], + [0, 1, -1, -1, -1, -1, -1], + ], + dtype=int, + ).reshape([1, self.nloc, sum(self.sel)]) + self.rcut = 2.2 + self.rcut_smth = 0.4 + # permutations + self.perm = np.array([2, 0, 1, 3], dtype=np.int32) + inv_perm = np.array([1, 2, 0, 3], dtype=np.int32) + # permute the coord and atype + self.coord_ext = np.concatenate( + [self.coord_ext, self.coord_ext[:, self.perm, :]], axis=0 + ).reshape(self.nf, self.nall * 3) + self.atype_ext = np.concatenate( + [self.atype_ext, self.atype_ext[:, self.perm]], axis=0 + ) + self.mapping = np.concatenate( + [self.mapping, self.mapping[:, self.perm]], axis=0 + ) + + # permute the nlist + nlist1 = self.nlist[:, self.perm[: self.nloc], :] + mask = nlist1 == -1 + nlist1 = inv_perm[nlist1] + nlist1 = np.where(mask, -1, nlist1) + self.nlist = np.concatenate([self.nlist, nlist1], axis=0) + self.atol = 1e-12 + + +class TestCaseSingleFrameWithNlistWithVirtual: + def setUp(self): + # nloc == 3, nall == 4 + self.nloc = 4 + self.nall = 5 + self.nf, self.nt = 2, 2 + self.coord_ext = np.array( + [ + [0, 0, 0], + [0, 0, 0], + [0, 1, 0], + [0, 0, 1], + [0, -2, 0], + ], + dtype=np.float64, + ).reshape([1, self.nall, 3]) + self.atype_ext = np.array([0, -1, 0, 1, 0], dtype=int).reshape([1, self.nall]) + # sel = [5, 2] + self.sel = [5, 2] + self.sel_mix = [7] + self.natoms = [3, 3, 2, 1] + self.nlist = np.array( + [ + [2, 4, -1, -1, -1, 3, -1], + [-1, -1, -1, -1, -1, -1, -1], + [0, -1, -1, -1, -1, 3, -1], + [0, 2, -1, -1, -1, -1, -1], + ], + dtype=int, + ).reshape([1, self.nloc, sum(self.sel)]) + self.rcut = 2.2 + self.rcut_smth = 0.4 + # permutations + self.perm = np.array([3, 0, 1, 2, 4], dtype=np.int32) + inv_perm = np.argsort(self.perm) + # permute the coord and atype + self.coord_ext = np.concatenate( + [self.coord_ext, self.coord_ext[:, self.perm, :]], axis=0 + ).reshape(self.nf, self.nall * 3) + self.atype_ext = np.concatenate( + [self.atype_ext, self.atype_ext[:, self.perm]], axis=0 + ) + # permute the nlist + nlist1 = self.nlist[:, self.perm[: self.nloc], :] + mask = nlist1 == -1 + nlist1 = inv_perm[nlist1] + nlist1 = np.where(mask, -1, nlist1) + self.nlist = np.concatenate([self.nlist, nlist1], axis=0) + self.get_real_mapping = np.array([[0, 2, 3], [0, 1, 3]], dtype=np.int32) + self.atol = 1e-12 + + +class TestCaseSingleFrameWithoutNlist: + def setUp(self): + # nloc == 3, nall == 4 + self.nloc = 3 + self.nf, self.nt = 1, 2 + self.coord = np.array( + [ + [0, 0, 0], + [0, 1, 0], + [0, 0, 1], + ], + dtype=np.float64, + ).reshape([1, self.nloc * 3]) + self.atype = np.array([0, 0, 1], dtype=int).reshape([1, self.nloc]) + self.cell = 2.0 * np.eye(3).reshape([1, 9]) + # sel = [5, 2] + self.sel = [16, 8] + self.sel_mix = [24] + self.natoms = [3, 3, 2, 1] + self.rcut = 2.2 + self.rcut_smth = 0.4 + self.atol = 1e-12 diff --git a/source/tests/universal/common/cases/descriptor/__init__.py b/source/tests/universal/common/cases/descriptor/__init__.py new file mode 100644 index 0000000000..6ceb116d85 --- /dev/null +++ b/source/tests/universal/common/cases/descriptor/__init__.py @@ -0,0 +1 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later diff --git a/source/tests/universal/common/cases/descriptor/descriptor.py b/source/tests/universal/common/cases/descriptor/descriptor.py new file mode 100644 index 0000000000..2f3bdbb1ee --- /dev/null +++ b/source/tests/universal/common/cases/descriptor/descriptor.py @@ -0,0 +1,11 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later + + +from .utils import ( + DescriptorTestCase, +) + + +class DescriptorTest(DescriptorTestCase): + def setUp(self) -> None: + DescriptorTestCase.setUp(self) diff --git a/source/tests/universal/common/cases/descriptor/utils.py b/source/tests/universal/common/cases/descriptor/utils.py new file mode 100644 index 0000000000..b49af0fdb6 --- /dev/null +++ b/source/tests/universal/common/cases/descriptor/utils.py @@ -0,0 +1,95 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +from copy import ( + deepcopy, +) + +import numpy as np + +from deepmd.dpmodel.utils import ( + PairExcludeMask, +) + +from ..cases import ( + TestCaseSingleFrameWithNlist, +) + + +class DescriptorTestCase(TestCaseSingleFrameWithNlist): + """Common test case for atomic model.""" + + def setUp(self): + TestCaseSingleFrameWithNlist.setUp(self) + self.input_dict = { + "ntypes": self.nt, + "rcut": self.rcut, + "rcut_smth": self.rcut_smth, + "sel": self.sel, + } + + def test_forward(self): + ret = [] + for module in self.modules_to_test: + module = self.forward_wrapper(module) + ret.append( + module( + self.coord_ext, + self.atype_ext, + self.nlist, + mapping=self.mapping, + ) + ) + for kk, vv in enumerate(ret[0]): + subret = [] + for rr in ret: + if rr is not None: + subret.append(rr[kk]) + if len(subret): + for ii, rr in enumerate(subret[1:]): + if subret[0] is None: + assert rr is None + else: + np.testing.assert_allclose( + subret[0], + rr, + err_msg=f"compare {kk} output between 0 and {ii}", + ) + + def test_exclude_types( + self, + ): + coord_ext_device = self.coord_ext + atype_ext_device = self.atype_ext + nlist_device = self.nlist + mapping_device = self.mapping + dd = self.forward_wrapper(self.module) + # only equal when set_davg_zero is True + serialize_dict = self.module.serialize() + + for em in [[[0, 1]], [[1, 1]]]: + ex_pair = PairExcludeMask(self.nt, em) + pair_mask = ex_pair.build_type_exclude_mask(nlist_device, atype_ext_device) + # exclude neighbors in the nlist + nlist_exclude = np.where(pair_mask == 1, nlist_device, -1) + rd_ex, _, _, _, sw_ex = dd( + coord_ext_device, + atype_ext_device, + nlist_exclude, + mapping=mapping_device, + ) + + # normal nlist but use exclude_types params + serialize_dict_em = deepcopy(serialize_dict) + if "list" not in serialize_dict_em: + serialize_dict_em.update({"exclude_types": em}) + else: + # for hybrid + for sd in serialize_dict_em["list"]: + sd.update({"exclude_types": em}) + dd0 = self.forward_wrapper(self.module.deserialize(serialize_dict_em)) + rd0, _, _, _, sw0 = dd0( + coord_ext_device, + atype_ext_device, + nlist_device, + mapping=mapping_device, + ) + np.testing.assert_allclose(rd0, rd_ex) diff --git a/source/tests/universal/dpmodel/descriptor/__init__.py b/source/tests/universal/dpmodel/descriptor/__init__.py new file mode 100644 index 0000000000..6ceb116d85 --- /dev/null +++ b/source/tests/universal/dpmodel/descriptor/__init__.py @@ -0,0 +1 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later diff --git a/source/tests/universal/dpmodel/descriptor/test_descriptor.py b/source/tests/universal/dpmodel/descriptor/test_descriptor.py new file mode 100644 index 0000000000..9d0253c54c --- /dev/null +++ b/source/tests/universal/dpmodel/descriptor/test_descriptor.py @@ -0,0 +1,84 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import unittest + +from deepmd.dpmodel.descriptor import ( + DescrptDPA1, + DescrptDPA2, + DescrptHybrid, + DescrptSeA, + DescrptSeR, + DescrptSeT, +) + +from ...common.cases.descriptor.descriptor import ( + DescriptorTest, +) +from ..backend import ( + DPTestCase, +) + + +class TestDescriptorSeADP(unittest.TestCase, DescriptorTest, DPTestCase): + def setUp(self): + DescriptorTest.setUp(self) + self.module = DescrptSeA(**self.input_dict) + + +class TestDescriptorSeRDP(unittest.TestCase, DescriptorTest, DPTestCase): + def setUp(self): + DescriptorTest.setUp(self) + self.module = DescrptSeR(**self.input_dict) + + +class TestDescriptorSeTDP(unittest.TestCase, DescriptorTest, DPTestCase): + def setUp(self): + DescriptorTest.setUp(self) + self.module = DescrptSeT(**self.input_dict) + + +class TestDescriptorDPA1DP(unittest.TestCase, DescriptorTest, DPTestCase): + def setUp(self): + DescriptorTest.setUp(self) + self.module = DescrptDPA1(**self.input_dict) + + +class TestDescriptorDPA2DP(unittest.TestCase, DescriptorTest, DPTestCase): + def setUp(self): + DescriptorTest.setUp(self) + self.input_dict = { + "ntypes": self.nt, + "repinit": { + "rcut": self.rcut, + "rcut_smth": self.rcut_smth, + "nsel": self.sel_mix, + }, + "repformer": { + "rcut": self.rcut / 2, + "rcut_smth": self.rcut_smth, + "nsel": self.sel_mix[0] // 2, + }, + } + self.module = DescrptDPA2(**self.input_dict) + + +class TestDescriptorHybridDP(unittest.TestCase, DescriptorTest, DPTestCase): + def setUp(self): + DescriptorTest.setUp(self) + ddsub0 = { + "type": "se_e2_a", + "ntypes": self.nt, + "rcut": self.rcut, + "rcut_smth": self.rcut_smth, + "sel": self.sel, + } + ddsub1 = { + "type": "dpa1", + "ntypes": self.nt, + "rcut": self.rcut, + "rcut_smth": self.rcut_smth, + "sel": self.sel_mix, + } + self.input_dict = { + "list": [ddsub0, ddsub1], + } + self.module = DescrptHybrid(**self.input_dict) diff --git a/source/tests/universal/pt/descriptor/__init__.py b/source/tests/universal/pt/descriptor/__init__.py new file mode 100644 index 0000000000..6ceb116d85 --- /dev/null +++ b/source/tests/universal/pt/descriptor/__init__.py @@ -0,0 +1 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later diff --git a/source/tests/universal/pt/descriptor/test_descriptor.py b/source/tests/universal/pt/descriptor/test_descriptor.py new file mode 100644 index 0000000000..87107a2f90 --- /dev/null +++ b/source/tests/universal/pt/descriptor/test_descriptor.py @@ -0,0 +1,84 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import unittest + +from deepmd.pt.model.descriptor import ( + DescrptDPA1, + DescrptDPA2, + DescrptHybrid, + DescrptSeA, + DescrptSeR, + DescrptSeT, +) + +from ...common.cases.descriptor.descriptor import ( + DescriptorTest, +) +from ..backend import ( + PTTestCase, +) + + +class TestDescriptorSeAPT(unittest.TestCase, DescriptorTest, PTTestCase): + def setUp(self): + DescriptorTest.setUp(self) + self.module = DescrptSeA(**self.input_dict) + + +class TestDescriptorSeRPT(unittest.TestCase, DescriptorTest, PTTestCase): + def setUp(self): + DescriptorTest.setUp(self) + self.module = DescrptSeR(**self.input_dict) + + +class TestDescriptorSeTPT(unittest.TestCase, DescriptorTest, PTTestCase): + def setUp(self): + DescriptorTest.setUp(self) + self.module = DescrptSeT(**self.input_dict) + + +class TestDescriptorDPA1PT(unittest.TestCase, DescriptorTest, PTTestCase): + def setUp(self): + DescriptorTest.setUp(self) + self.module = DescrptDPA1(**self.input_dict) + + +class TestDescriptorDPA2PT(unittest.TestCase, DescriptorTest, PTTestCase): + def setUp(self): + DescriptorTest.setUp(self) + self.input_dict = { + "ntypes": self.nt, + "repinit": { + "rcut": self.rcut, + "rcut_smth": self.rcut_smth, + "nsel": self.sel_mix, + }, + "repformer": { + "rcut": self.rcut / 2, + "rcut_smth": self.rcut_smth, + "nsel": self.sel_mix[0] // 2, + }, + } + self.module = DescrptDPA2(**self.input_dict) + + +class TestDescriptorHybridPT(unittest.TestCase, DescriptorTest, PTTestCase): + def setUp(self): + DescriptorTest.setUp(self) + ddsub0 = { + "type": "se_e2_a", + "ntypes": self.nt, + "rcut": self.rcut, + "rcut_smth": self.rcut_smth, + "sel": self.sel, + } + ddsub1 = { + "type": "dpa1", + "ntypes": self.nt, + "rcut": self.rcut, + "rcut_smth": self.rcut_smth, + "sel": self.sel_mix, + } + self.input_dict = { + "list": [ddsub0, ddsub1], + } + self.module = DescrptHybrid(**self.input_dict) From 271c3cbd6566fb36b7c21ea489832336d8a3b23e Mon Sep 17 00:00:00 2001 From: Duo <50307526+iProzd@users.noreply.github.com> Date: Fri, 31 May 2024 17:34:09 +0800 Subject: [PATCH 02/10] Delete test_exclude_types.py --- source/tests/pt/model/test_exclude_types.py | 151 -------------------- 1 file changed, 151 deletions(-) delete mode 100644 source/tests/pt/model/test_exclude_types.py diff --git a/source/tests/pt/model/test_exclude_types.py b/source/tests/pt/model/test_exclude_types.py deleted file mode 100644 index cf7eb55705..0000000000 --- a/source/tests/pt/model/test_exclude_types.py +++ /dev/null @@ -1,151 +0,0 @@ -# SPDX-License-Identifier: LGPL-3.0-or-later -import unittest - -import numpy as np -import torch - -from deepmd.pt.model.descriptor import ( - DescrptDPA1, - DescrptDPA2, - DescrptHybrid, - DescrptSeA, - DescrptSeR, - DescrptSeT, -) -from deepmd.pt.utils import ( - PairExcludeMask, - env, -) -from deepmd.pt.utils.env import ( - PRECISION_DICT, -) - -from .test_env_mat import ( - TestCaseSingleFrameWithNlist, -) -from .test_mlp import ( - get_tols, -) - -dtype = env.GLOBAL_PT_FLOAT_PRECISION - - -class DescrptExcludeTypes(TestCaseSingleFrameWithNlist): - def setUp(self): - TestCaseSingleFrameWithNlist.setUp(self) - self.input_dict = { - "ntypes": self.nt, - "rcut": self.rcut, - "rcut_smth": self.rcut_smth, - "sel": self.sel, - } - - def test_exclude_types( - self, - ): - dtype = PRECISION_DICT["float64"] - rtol, atol = get_tols("float64") - coord_ext_device = torch.tensor(self.coord_ext, dtype=dtype, device=env.DEVICE) - atype_ext_device = torch.tensor(self.atype_ext, dtype=int, device=env.DEVICE) - nlist_device = torch.tensor(self.nlist, dtype=int, device=env.DEVICE) - mapping_device = torch.tensor(self.mapping, dtype=int, device=env.DEVICE) - dd = self.descrpt(**self.input_dict).to(env.DEVICE) - - for em in [[[0, 1]], [[1, 1]]]: - dd0 = self.descrpt(**self.input_dict, exclude_types=em).to(env.DEVICE) - # only equal when set_davg_zero is True - # dd0.sea.mean = torch.tensor(davg, dtype=dtype, device=env.DEVICE) - # dd0.sea.dstd = torch.tensor(dstd, dtype=dtype, device=env.DEVICE) - dd.load_state_dict(dd0.state_dict()) - - ex_pair = PairExcludeMask(self.nt, em) - pair_mask = ex_pair(nlist_device, atype_ext_device) - # exclude neighbors in the nlist - nlist_exclude = torch.where(pair_mask == 1, nlist_device, -1) - - rd0, _, _, _, sw0 = dd0( - coord_ext_device, - atype_ext_device, - nlist_device, - mapping=mapping_device, - ) - - rd_ex, _, _, _, sw_ex = dd( - coord_ext_device, - atype_ext_device, - nlist_exclude, - mapping=mapping_device, - ) - - np.testing.assert_allclose( - rd0.detach().cpu().numpy(), - rd_ex.detach().cpu().numpy(), - rtol=rtol, - atol=atol, - ) - - -class TestDescrptExcludeTypesSeA(DescrptExcludeTypes, unittest.TestCase): - def setUp(self): - super().setUp() - self.descrpt = DescrptSeA - - -class TestDescrptExcludeTypesSeR(DescrptExcludeTypes, unittest.TestCase): - def setUp(self): - super().setUp() - self.descrpt = DescrptSeR - - -class TestDescrptExcludeTypesSeT(DescrptExcludeTypes, unittest.TestCase): - def setUp(self): - super().setUp() - self.descrpt = DescrptSeT - - -class TestDescrptExcludeTypesDPA1(DescrptExcludeTypes, unittest.TestCase): - def setUp(self): - super().setUp() - self.descrpt = DescrptDPA1 - - -class TestDescrptExcludeTypesDPA2(DescrptExcludeTypes, unittest.TestCase): - def setUp(self): - super().setUp() - self.descrpt = DescrptDPA2 - self.input_dict = { - "ntypes": self.nt, - "repinit": { - "rcut": self.rcut, - "rcut_smth": self.rcut_smth, - "nsel": self.sel_mix, - }, - "repformer": { - "rcut": self.rcut / 2, - "rcut_smth": self.rcut_smth, - "nsel": self.sel_mix[0] // 2, - }, - } - - -class TestDescrptExcludeTypesHybrid(DescrptExcludeTypes, unittest.TestCase): - def setUp(self): - super().setUp() - self.descrpt = DescrptHybrid - ddsub0 = { - "type": "se_e2_a", - "ntypes": self.nt, - "rcut": self.rcut, - "rcut_smth": self.rcut_smth, - "sel": self.sel, - } - ddsub1 = { - "type": "dpa1", - "ntypes": self.nt, - "rcut": self.rcut, - "rcut_smth": self.rcut_smth, - "sel": self.sel_mix, - } - self.input_dict = { - "list": [ddsub0, ddsub1], - } From dc83232088244ce8c16ade508a7a901df3b789f3 Mon Sep 17 00:00:00 2001 From: Duo <50307526+iProzd@users.noreply.github.com> Date: Fri, 31 May 2024 17:57:32 +0800 Subject: [PATCH 03/10] Update se_atten.py --- deepmd/pt/model/descriptor/se_atten.py | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/deepmd/pt/model/descriptor/se_atten.py b/deepmd/pt/model/descriptor/se_atten.py index d26578dc74..84437ab911 100644 --- a/deepmd/pt/model/descriptor/se_atten.py +++ b/deepmd/pt/model/descriptor/se_atten.py @@ -495,11 +495,7 @@ def forward( atype_tebd_nlist = atype_tebd_nlist.view(nb, nloc, nnei, nt) # nb x nloc x nnei exclude_mask = self.emask(nlist, extended_atype) - nlist_mask = torch.where( - exclude_mask != 0, - nlist_mask, - torch.zeros_like(nlist_mask, dtype=torch.bool), - ) + nlist_mask &= exclude_mask != 0 # beyond the cutoff sw should be 0.0 sw = sw.masked_fill(~nlist_mask, 0.0) # (nb x nloc) x nnei From dfa4944b1f802bab8cc6ccefd9aa00cfd875477e Mon Sep 17 00:00:00 2001 From: Duo <50307526+iProzd@users.noreply.github.com> Date: Fri, 31 May 2024 18:20:24 +0800 Subject: [PATCH 04/10] Update utils.py --- source/tests/universal/common/cases/descriptor/utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/source/tests/universal/common/cases/descriptor/utils.py b/source/tests/universal/common/cases/descriptor/utils.py index b49af0fdb6..d8bde17dc7 100644 --- a/source/tests/universal/common/cases/descriptor/utils.py +++ b/source/tests/universal/common/cases/descriptor/utils.py @@ -15,7 +15,7 @@ class DescriptorTestCase(TestCaseSingleFrameWithNlist): - """Common test case for atomic model.""" + """Common test case for descriptor.""" def setUp(self): TestCaseSingleFrameWithNlist.setUp(self) From 3f178987bd3dd2df812d9dc4c1e6bd3acfb543e4 Mon Sep 17 00:00:00 2001 From: Duo <50307526+iProzd@users.noreply.github.com> Date: Fri, 31 May 2024 18:36:58 +0800 Subject: [PATCH 05/10] Update cases.py --- source/tests/universal/common/cases/cases.py | 1 + 1 file changed, 1 insertion(+) diff --git a/source/tests/universal/common/cases/cases.py b/source/tests/universal/common/cases/cases.py index 797ab188a8..a8c9a7cd71 100644 --- a/source/tests/universal/common/cases/cases.py +++ b/source/tests/universal/common/cases/cases.py @@ -2,6 +2,7 @@ import numpy as np +# originally copied from source/tests/pt/model/test_env_mat.py class TestCaseSingleFrameWithNlist: def setUp(self): # nloc == 3, nall == 4 From 1955e6b1f41e9f559466d82d0eeeda47b0d22ce2 Mon Sep 17 00:00:00 2001 From: Duo <50307526+iProzd@users.noreply.github.com> Date: Fri, 31 May 2024 18:49:05 +0800 Subject: [PATCH 06/10] Update se_atten.py --- deepmd/pt/model/descriptor/se_atten.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/deepmd/pt/model/descriptor/se_atten.py b/deepmd/pt/model/descriptor/se_atten.py index 84437ab911..f5d732b1e7 100644 --- a/deepmd/pt/model/descriptor/se_atten.py +++ b/deepmd/pt/model/descriptor/se_atten.py @@ -495,7 +495,7 @@ def forward( atype_tebd_nlist = atype_tebd_nlist.view(nb, nloc, nnei, nt) # nb x nloc x nnei exclude_mask = self.emask(nlist, extended_atype) - nlist_mask &= exclude_mask != 0 + nlist_mask = nlist_mask & (exclude_mask != 0) # beyond the cutoff sw should be 0.0 sw = sw.masked_fill(~nlist_mask, 0.0) # (nb x nloc) x nnei From 545430ff3d8b5ca2bfe06b83ee6694f59c82725b Mon Sep 17 00:00:00 2001 From: Duo <50307526+iProzd@users.noreply.github.com> Date: Mon, 3 Jun 2024 14:53:19 +0800 Subject: [PATCH 07/10] resolve conversations --- deepmd/dpmodel/descriptor/dpa1.py | 3 +-- deepmd/pt/model/descriptor/se_atten.py | 6 +++--- deepmd/tf/descriptor/se_atten.py | 16 ++++++++-------- .../universal/common/cases/descriptor/utils.py | 2 +- 4 files changed, 13 insertions(+), 14 deletions(-) diff --git a/deepmd/dpmodel/descriptor/dpa1.py b/deepmd/dpmodel/descriptor/dpa1.py index d9f8570cd0..0f80406bb4 100644 --- a/deepmd/dpmodel/descriptor/dpa1.py +++ b/deepmd/dpmodel/descriptor/dpa1.py @@ -813,6 +813,7 @@ def call( exclude_mask = exclude_mask.reshape(nf * nloc, nnei) # nfnl x nnei nlist = nlist.reshape(nf * nloc, nnei) + nlist = np.where(exclude_mask, nlist, -1) # nfnl x nnei x 4 dmatrix = dmatrix.reshape(nf * nloc, nnei, 4) # nfnl x nnei x 1 @@ -823,8 +824,6 @@ def call( atype_embd_nnei = np.tile(atype_embd[:, np.newaxis, :], (1, nnei, 1)) # nfnl x nnei nlist_mask = nlist != -1 - # nfnl x nnei - nlist_mask = nlist_mask * exclude_mask.astype(bool) # nfnl x nnei x 1 sw = np.where(nlist_mask[:, :, None], sw, 0.0) nlist_masked = np.where(nlist_mask, nlist, 0) diff --git a/deepmd/pt/model/descriptor/se_atten.py b/deepmd/pt/model/descriptor/se_atten.py index f5d732b1e7..181277b637 100644 --- a/deepmd/pt/model/descriptor/se_atten.py +++ b/deepmd/pt/model/descriptor/se_atten.py @@ -478,6 +478,9 @@ def forward( self.rcut_smth, protection=self.env_protection, ) + # nb x nloc x nnei + exclude_mask = self.emask(nlist, extended_atype) + nlist = torch.where(exclude_mask != 0, nlist, -1) nlist_mask = nlist != -1 nlist = torch.where(nlist == -1, 0, nlist) sw = torch.squeeze(sw, -1) @@ -493,9 +496,6 @@ def forward( atype_tebd_nlist = torch.gather(atype_tebd_ext, dim=1, index=index) # nb x nloc x nnei x nt atype_tebd_nlist = atype_tebd_nlist.view(nb, nloc, nnei, nt) - # nb x nloc x nnei - exclude_mask = self.emask(nlist, extended_atype) - nlist_mask = nlist_mask & (exclude_mask != 0) # beyond the cutoff sw should be 0.0 sw = sw.masked_fill(~nlist_mask, 0.0) # (nb x nloc) x nnei diff --git a/deepmd/tf/descriptor/se_atten.py b/deepmd/tf/descriptor/se_atten.py index c05d31ba60..b240f00647 100644 --- a/deepmd/tf/descriptor/se_atten.py +++ b/deepmd/tf/descriptor/se_atten.py @@ -718,18 +718,12 @@ def _pass_filter( tf.shape(inputs_i)[0], self.nei_type_vec, # extra input for atten ) - # (nsamples * natoms * nnei, 1) + # (nframes * nloc * nnei, 1) nei_exclude_mask = tf.slice( tf.reshape(tf.cast(mask, self.filter_precision), [-1, 4]), [0, 0], [-1, 1], ) - # (nsamples * natoms, 1, nnei) - self.nmask *= tf.reshape( - nei_exclude_mask, - [-1, 1, self.sel_all_a[0]], - ) - self.negative_mask = -(2 << 32) * (1.0 - self.nmask) if self.smooth: inputs_i = tf.where( tf.cast(mask, tf.bool), @@ -739,12 +733,18 @@ def _pass_filter( tf.reshape(self.avg_looked_up, [-1, 1]), [1, self.ndescrpt] ), ) - # (nsamples, natoms, nnei) + # (nframes, nloc, nnei) self.recovered_switch *= tf.reshape( nei_exclude_mask, [-1, natoms[0], self.sel_all_a[0]], ) else: + # (nframes * nloc, 1, nnei) + self.nmask *= tf.reshape( + nei_exclude_mask, + [-1, 1, self.sel_all_a[0]], + ) + self.negative_mask = -(2 << 32) * (1.0 - self.nmask) inputs_i *= mask if nvnmd_cfg.enable and nvnmd_cfg.quantize_descriptor: inputs_i = descrpt2r4(inputs_i, atype) diff --git a/source/tests/universal/common/cases/descriptor/utils.py b/source/tests/universal/common/cases/descriptor/utils.py index d8bde17dc7..aa1a8c21d4 100644 --- a/source/tests/universal/common/cases/descriptor/utils.py +++ b/source/tests/universal/common/cases/descriptor/utils.py @@ -26,7 +26,7 @@ def setUp(self): "sel": self.sel, } - def test_forward(self): + def test_forward_consistency(self): ret = [] for module in self.modules_to_test: module = self.forward_wrapper(module) From c7aaef504c4bad06b94329dbb53c3971e760c365 Mon Sep 17 00:00:00 2001 From: Duo <50307526+iProzd@users.noreply.github.com> Date: Tue, 4 Jun 2024 18:17:56 +0800 Subject: [PATCH 08/10] Update test_cuda.yml --- .github/workflows/test_cuda.yml | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/.github/workflows/test_cuda.yml b/.github/workflows/test_cuda.yml index d7ff72bafa..f20e86860a 100644 --- a/.github/workflows/test_cuda.yml +++ b/.github/workflows/test_cuda.yml @@ -34,10 +34,7 @@ jobs: uses: mpi4py/setup-mpi@v1 with: mpi: mpich - - uses: lukka/get-cmake@latest - with: - useLocalCache: true - useCloudCache: false + - run: pip install cmake - name: Install wget and unzip run: apt-get update && apt-get install -y wget unzip - run: | From fbddb4dd59e0c5793535678a35c02e990675615f Mon Sep 17 00:00:00 2001 From: Duo <50307526+iProzd@users.noreply.github.com> Date: Tue, 4 Jun 2024 22:52:04 +0800 Subject: [PATCH 09/10] Revert "Update test_cuda.yml" This reverts commit c7aaef504c4bad06b94329dbb53c3971e760c365. --- .github/workflows/test_cuda.yml | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/.github/workflows/test_cuda.yml b/.github/workflows/test_cuda.yml index f20e86860a..d7ff72bafa 100644 --- a/.github/workflows/test_cuda.yml +++ b/.github/workflows/test_cuda.yml @@ -34,7 +34,10 @@ jobs: uses: mpi4py/setup-mpi@v1 with: mpi: mpich - - run: pip install cmake + - uses: lukka/get-cmake@latest + with: + useLocalCache: true + useCloudCache: false - name: Install wget and unzip run: apt-get update && apt-get install -y wget unzip - run: | From 0e11a4921a59b2d4974b570409c8ae71f6abb90c Mon Sep 17 00:00:00 2001 From: Duo <50307526+iProzd@users.noreply.github.com> Date: Tue, 4 Jun 2024 22:52:44 +0800 Subject: [PATCH 10/10] Update test_cuda.yml --- .github/workflows/test_cuda.yml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/.github/workflows/test_cuda.yml b/.github/workflows/test_cuda.yml index d7ff72bafa..d97b1f9431 100644 --- a/.github/workflows/test_cuda.yml +++ b/.github/workflows/test_cuda.yml @@ -34,12 +34,12 @@ jobs: uses: mpi4py/setup-mpi@v1 with: mpi: mpich + - name: Install wget and unzip + run: apt-get update && apt-get install -y wget unzip - uses: lukka/get-cmake@latest with: useLocalCache: true useCloudCache: false - - name: Install wget and unzip - run: apt-get update && apt-get install -y wget unzip - run: | wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/cuda-keyring_1.0-1_all.deb \ && sudo dpkg -i cuda-keyring_1.0-1_all.deb \