From df49f1af67fde931ccb7662dcb78e0e50a5fa988 Mon Sep 17 00:00:00 2001 From: anyangml Date: Wed, 8 May 2024 05:04:38 +0000 Subject: [PATCH 1/8] fix: raise error for non mixed type model --- deepmd/dpmodel/atomic_model/linear_atomic_model.py | 9 +++++++++ deepmd/pt/model/atomic_model/linear_atomic_model.py | 9 +++++++++ 2 files changed, 18 insertions(+) diff --git a/deepmd/dpmodel/atomic_model/linear_atomic_model.py b/deepmd/dpmodel/atomic_model/linear_atomic_model.py index b38d309fd7..2e6a95b25a 100644 --- a/deepmd/dpmodel/atomic_model/linear_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/linear_atomic_model.py @@ -54,6 +54,15 @@ def __init__( ): super().__init__(type_map, **kwargs) super().init_out_stat() + + # check all sub models are of mixed type. + model_mixed_type = [] + for m in models: + if not m.mixed_types(): + model_mixed_type.append(m) + if len(model_mixed_type) > 0: + raise ValueError(f"LinearAtomicModel only supports AtomicModel of mixed type, the following models are not mixed type: {err_msg}.") + self.models = models sub_model_type_maps = [md.get_type_map() for md in models] err_msg = [] diff --git a/deepmd/pt/model/atomic_model/linear_atomic_model.py b/deepmd/pt/model/atomic_model/linear_atomic_model.py index bf03a68f31..72b0cca403 100644 --- a/deepmd/pt/model/atomic_model/linear_atomic_model.py +++ b/deepmd/pt/model/atomic_model/linear_atomic_model.py @@ -61,6 +61,15 @@ def __init__( ): super().__init__(type_map, **kwargs) super().init_out_stat() + + # check all sub models are of mixed type. + model_mixed_type = [] + for m in models: + if not m.mixed_types(): + model_mixed_type.append(m) + if len(model_mixed_type) > 0: + raise ValueError(f"LinearAtomicModel only supports AtomicModel of mixed type, the following models are not mixed type: {err_msg}.") + self.models = torch.nn.ModuleList(models) sub_model_type_maps = [md.get_type_map() for md in models] err_msg = [] From 3a0be54fed81970a13428a91650352be00b0c6b4 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Wed, 8 May 2024 05:11:01 +0000 Subject: [PATCH 2/8] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- deepmd/dpmodel/atomic_model/linear_atomic_model.py | 4 +++- deepmd/pt/model/atomic_model/linear_atomic_model.py | 4 +++- 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/deepmd/dpmodel/atomic_model/linear_atomic_model.py b/deepmd/dpmodel/atomic_model/linear_atomic_model.py index 2e6a95b25a..e39bd92be7 100644 --- a/deepmd/dpmodel/atomic_model/linear_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/linear_atomic_model.py @@ -61,7 +61,9 @@ def __init__( if not m.mixed_types(): model_mixed_type.append(m) if len(model_mixed_type) > 0: - raise ValueError(f"LinearAtomicModel only supports AtomicModel of mixed type, the following models are not mixed type: {err_msg}.") + raise ValueError( + f"LinearAtomicModel only supports AtomicModel of mixed type, the following models are not mixed type: {err_msg}." + ) self.models = models sub_model_type_maps = [md.get_type_map() for md in models] diff --git a/deepmd/pt/model/atomic_model/linear_atomic_model.py b/deepmd/pt/model/atomic_model/linear_atomic_model.py index 72b0cca403..eeda3944e5 100644 --- a/deepmd/pt/model/atomic_model/linear_atomic_model.py +++ b/deepmd/pt/model/atomic_model/linear_atomic_model.py @@ -68,7 +68,9 @@ def __init__( if not m.mixed_types(): model_mixed_type.append(m) if len(model_mixed_type) > 0: - raise ValueError(f"LinearAtomicModel only supports AtomicModel of mixed type, the following models are not mixed type: {err_msg}.") + raise ValueError( + f"LinearAtomicModel only supports AtomicModel of mixed type, the following models are not mixed type: {err_msg}." + ) self.models = torch.nn.ModuleList(models) sub_model_type_maps = [md.get_type_map() for md in models] From 90da1ad09e9e1ed510484f3d173358bfd1af8765 Mon Sep 17 00:00:00 2001 From: anyangml Date: Wed, 8 May 2024 05:17:45 +0000 Subject: [PATCH 3/8] fix: precommit --- deepmd/dpmodel/atomic_model/linear_atomic_model.py | 2 +- deepmd/pt/model/atomic_model/linear_atomic_model.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/deepmd/dpmodel/atomic_model/linear_atomic_model.py b/deepmd/dpmodel/atomic_model/linear_atomic_model.py index e39bd92be7..7dff9078c5 100644 --- a/deepmd/dpmodel/atomic_model/linear_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/linear_atomic_model.py @@ -62,7 +62,7 @@ def __init__( model_mixed_type.append(m) if len(model_mixed_type) > 0: raise ValueError( - f"LinearAtomicModel only supports AtomicModel of mixed type, the following models are not mixed type: {err_msg}." + f"LinearAtomicModel only supports AtomicModel of mixed type, the following models are not mixed type: {model_mixed_type}." ) self.models = models diff --git a/deepmd/pt/model/atomic_model/linear_atomic_model.py b/deepmd/pt/model/atomic_model/linear_atomic_model.py index eeda3944e5..d21c65c257 100644 --- a/deepmd/pt/model/atomic_model/linear_atomic_model.py +++ b/deepmd/pt/model/atomic_model/linear_atomic_model.py @@ -69,7 +69,7 @@ def __init__( model_mixed_type.append(m) if len(model_mixed_type) > 0: raise ValueError( - f"LinearAtomicModel only supports AtomicModel of mixed type, the following models are not mixed type: {err_msg}." + f"LinearAtomicModel only supports AtomicModel of mixed type, the following models are not mixed type: {model_mixed_type}." ) self.models = torch.nn.ModuleList(models) From 7dabcf548dac30e28825f17615297e8f6f3e9849 Mon Sep 17 00:00:00 2001 From: anyangml Date: Wed, 8 May 2024 08:07:19 +0000 Subject: [PATCH 4/8] fix: nlist copy --- deepmd/pt/model/descriptor/se_atten.py | 7 ++++--- .../pt/model/test_linear_atomic_model.py | 12 ++++++----- source/tests/pt/model/test_permutation.py | 20 +++++++++++++------ 3 files changed, 25 insertions(+), 14 deletions(-) diff --git a/deepmd/pt/model/descriptor/se_atten.py b/deepmd/pt/model/descriptor/se_atten.py index 9bf4788bf2..07d36e595a 100644 --- a/deepmd/pt/model/descriptor/se_atten.py +++ b/deepmd/pt/model/descriptor/se_atten.py @@ -464,7 +464,8 @@ def forward( protection=self.env_protection, ) nlist_mask = nlist != -1 - nlist[nlist == -1] = 0 + nlist_copy = nlist.detach().clone() + nlist_copy[nlist == -1] = 0 sw = torch.squeeze(sw, -1) # beyond the cutoff sw should be 0.0 sw = sw.masked_fill(~nlist_mask, 0.0) @@ -475,13 +476,13 @@ def forward( nt = extended_atype_embd.shape[-1] atype_tebd_ext = extended_atype_embd # nb x (nloc x nnei) x nt - index = nlist.reshape(nb, nloc * nnei).unsqueeze(-1).expand(-1, -1, nt) + index = nlist_copy.reshape(nb, nloc * nnei).unsqueeze(-1).expand(-1, -1, nt) # nb x (nloc x nnei) x nt 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).view(nb * nloc, nnei) + exclude_mask = self.emask(nlist_copy, extended_atype).view(nb * nloc, nnei) if self.old_impl: assert self.filter_layers_old is not None dmatrix = dmatrix.view( diff --git a/source/tests/pt/model/test_linear_atomic_model.py b/source/tests/pt/model/test_linear_atomic_model.py index 7f24ffdc53..7104095250 100644 --- a/source/tests/pt/model/test_linear_atomic_model.py +++ b/source/tests/pt/model/test_linear_atomic_model.py @@ -15,8 +15,8 @@ DPZBLLinearEnergyAtomicModel, PairTabAtomicModel, ) -from deepmd.pt.model.descriptor.se_a import ( - DescrptSeA, +from deepmd.pt.model.descriptor import ( + DescrptDPA1, ) from deepmd.pt.model.model import ( DPZBLModel, @@ -55,10 +55,11 @@ def test_pairwise(self, mock_loadtxt): extended_atype = torch.tensor([[0, 0]], device=env.DEVICE) nlist = torch.tensor([[[1], [-1]]], device=env.DEVICE) - ds = DescrptSeA( + ds = DescrptDPA1( rcut_smth=0.3, rcut=0.4, sel=[3], + ntypes=2, ).to(env.DEVICE) ft = InvarFitting( "energy", @@ -128,10 +129,11 @@ def setUp(self, mock_loadtxt): [0.02, 0.25, 0.4, 0.75], ] ) - ds = DescrptSeA( + ds = DescrptDPA1( self.rcut, self.rcut_smth, - self.sel, + sum(self.sel), + self.nt, ).to(env.DEVICE) ft = InvarFitting( "energy", diff --git a/source/tests/pt/model/test_permutation.py b/source/tests/pt/model/test_permutation.py index b4cd133200..7e30d761e2 100644 --- a/source/tests/pt/model/test_permutation.py +++ b/source/tests/pt/model/test_permutation.py @@ -68,14 +68,22 @@ "sw_rmin": 0.2, "sw_rmax": 1.0, "descriptor": { - "type": "se_e2_a", - "sel": [46, 92, 4], - "rcut_smth": 0.50, - "rcut": 6.00, + "type": "se_atten", + "sel": 40, + "rcut_smth": 0.5, + "rcut": 4.0, "neuron": [25, 50, 100], - "resnet_dt": False, "axis_neuron": 16, - "seed": 1, + "attn": 64, + "attn_layer": 2, + "attn_dotr": True, + "attn_mask": False, + "activation_function": "tanh", + "scaling_factor": 1.0, + "normalize": False, + "temperature": 1.0, + "set_davg_zero": True, + "type_one_side": True, }, "fitting_net": { "neuron": [24, 24, 24], From 88fedda1a3092719d3ad9b21d089b838da9f4615 Mon Sep 17 00:00:00 2001 From: anyangml Date: Wed, 8 May 2024 08:35:06 +0000 Subject: [PATCH 5/8] fix: numpy --- .../tests/common/dpmodel/test_linear_atomic_model.py | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/source/tests/common/dpmodel/test_linear_atomic_model.py b/source/tests/common/dpmodel/test_linear_atomic_model.py index 832d1de106..ab7382e019 100644 --- a/source/tests/common/dpmodel/test_linear_atomic_model.py +++ b/source/tests/common/dpmodel/test_linear_atomic_model.py @@ -15,8 +15,8 @@ from deepmd.dpmodel.atomic_model.pairtab_atomic_model import ( PairTabAtomicModel, ) -from deepmd.dpmodel.descriptor.se_e2_a import ( - DescrptSeA, +from deepmd.dpmodel.descriptor import ( + DescrptDPA1, ) from deepmd.dpmodel.fitting.invar_fitting import ( InvarFitting, @@ -39,10 +39,11 @@ def test_pairwise(self, mock_loadtxt): extended_atype = np.array([[0, 0]]) nlist = np.array([[[1], [-1]]]) - ds = DescrptSeA( + ds = DescrptDPA1( rcut_smth=0.3, rcut=0.4, sel=[3], + ntypes=2, ) ft = InvarFitting( "energy", @@ -134,10 +135,11 @@ def setUp(self, mock_loadtxt): [0.02, 0.25, 0.4, 0.75], ] ) - ds = DescrptSeA( + ds = DescrptDPA1( self.rcut, self.rcut_smth, - self.sel, + sum(self.sel), + self.nt, ) ft = InvarFitting( "energy", From 1e7fd2f83259dea690be5549b35858157583e4a9 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Wed, 8 May 2024 08:35:56 +0000 Subject: [PATCH 6/8] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- source/tests/common/dpmodel/test_linear_atomic_model.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/source/tests/common/dpmodel/test_linear_atomic_model.py b/source/tests/common/dpmodel/test_linear_atomic_model.py index ab7382e019..b7bf310676 100644 --- a/source/tests/common/dpmodel/test_linear_atomic_model.py +++ b/source/tests/common/dpmodel/test_linear_atomic_model.py @@ -139,7 +139,7 @@ def setUp(self, mock_loadtxt): self.rcut, self.rcut_smth, sum(self.sel), - self.nt, + self.nt, ) ft = InvarFitting( "energy", From 53446dde68cbb4d48c337451e43da0a746681a42 Mon Sep 17 00:00:00 2001 From: Anyang Peng <137014849+anyangml@users.noreply.github.com> Date: Thu, 9 May 2024 09:46:12 +0800 Subject: [PATCH 7/8] Update deepmd/pt/model/descriptor/se_atten.py Co-authored-by: Han Wang <92130845+wanghan-iapcm@users.noreply.github.com> Signed-off-by: Anyang Peng <137014849+anyangml@users.noreply.github.com> --- deepmd/pt/model/descriptor/se_atten.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/deepmd/pt/model/descriptor/se_atten.py b/deepmd/pt/model/descriptor/se_atten.py index 07d36e595a..099fbc42a4 100644 --- a/deepmd/pt/model/descriptor/se_atten.py +++ b/deepmd/pt/model/descriptor/se_atten.py @@ -464,8 +464,7 @@ def forward( protection=self.env_protection, ) nlist_mask = nlist != -1 - nlist_copy = nlist.detach().clone() - nlist_copy[nlist == -1] = 0 + 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) From 8fea6f7c41408de91059e73b5c794685086d6e3c Mon Sep 17 00:00:00 2001 From: anyangml Date: Thu, 9 May 2024 01:49:13 +0000 Subject: [PATCH 8/8] fix: precommit --- deepmd/pt/model/descriptor/se_atten.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/deepmd/pt/model/descriptor/se_atten.py b/deepmd/pt/model/descriptor/se_atten.py index 099fbc42a4..d573ba9b7f 100644 --- a/deepmd/pt/model/descriptor/se_atten.py +++ b/deepmd/pt/model/descriptor/se_atten.py @@ -475,13 +475,13 @@ def forward( nt = extended_atype_embd.shape[-1] atype_tebd_ext = extended_atype_embd # nb x (nloc x nnei) x nt - index = nlist_copy.reshape(nb, nloc * nnei).unsqueeze(-1).expand(-1, -1, nt) + index = nlist.reshape(nb, nloc * nnei).unsqueeze(-1).expand(-1, -1, nt) # nb x (nloc x nnei) x nt 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_copy, extended_atype).view(nb * nloc, nnei) + exclude_mask = self.emask(nlist, extended_atype).view(nb * nloc, nnei) if self.old_impl: assert self.filter_layers_old is not None dmatrix = dmatrix.view(