From 3b59e2b1c45efad20c4f7777300028eabfb7d856 Mon Sep 17 00:00:00 2001 From: Anyang Peng <137014849+anyangml@users.noreply.github.com> Date: Thu, 9 May 2024 13:36:05 +0800 Subject: [PATCH] Fix: raise error for non mixed type model (#3759) ## Summary by CodeRabbit - **Bug Fixes** - Enhanced error handling in model initialization to ensure compatibility, now requiring all sub-models to be of mixed type. - **Refactor** - Modified code in `se_atten.py` to create a copy of `nlist` before modifying it, preserving the original `nlist` during subsequent operations. - **Tests** - Updated import statement and instantiation parameters in `test_linear_atomic_model.py` for `DescrptDPA1` class. - Altered parameters in the `descriptor` section of `test_permutation.py` for improved testing. --------- Signed-off-by: Anyang Peng <137014849+anyangml@users.noreply.github.com> Co-authored-by: anyangml Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Han Wang <92130845+wanghan-iapcm@users.noreply.github.com> --- .../atomic_model/linear_atomic_model.py | 11 ++++++++++ .../model/atomic_model/linear_atomic_model.py | 11 ++++++++++ deepmd/pt/model/descriptor/se_atten.py | 2 +- .../dpmodel/test_linear_atomic_model.py | 12 ++++++----- .../pt/model/test_linear_atomic_model.py | 12 ++++++----- source/tests/pt/model/test_permutation.py | 20 +++++++++++++------ 6 files changed, 51 insertions(+), 17 deletions(-) diff --git a/deepmd/dpmodel/atomic_model/linear_atomic_model.py b/deepmd/dpmodel/atomic_model/linear_atomic_model.py index b38d309fd7..7dff9078c5 100644 --- a/deepmd/dpmodel/atomic_model/linear_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/linear_atomic_model.py @@ -54,6 +54,17 @@ 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: {model_mixed_type}." + ) + 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..d21c65c257 100644 --- a/deepmd/pt/model/atomic_model/linear_atomic_model.py +++ b/deepmd/pt/model/atomic_model/linear_atomic_model.py @@ -61,6 +61,17 @@ 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: {model_mixed_type}." + ) + self.models = torch.nn.ModuleList(models) sub_model_type_maps = [md.get_type_map() for md in models] err_msg = [] diff --git a/deepmd/pt/model/descriptor/se_atten.py b/deepmd/pt/model/descriptor/se_atten.py index 9bf4788bf2..d573ba9b7f 100644 --- a/deepmd/pt/model/descriptor/se_atten.py +++ b/deepmd/pt/model/descriptor/se_atten.py @@ -464,7 +464,7 @@ def forward( protection=self.env_protection, ) nlist_mask = nlist != -1 - nlist[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) diff --git a/source/tests/common/dpmodel/test_linear_atomic_model.py b/source/tests/common/dpmodel/test_linear_atomic_model.py index 832d1de106..b7bf310676 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", 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 6159de199a..a54d3159e2 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],