From 96940f9430754fd414e471d46a46538fbdda69c1 Mon Sep 17 00:00:00 2001 From: Han Wang Date: Sat, 27 Jan 2024 22:57:21 +0800 Subject: [PATCH 01/14] unify the output of descriptors. --- deepmd/model_format/se_e2_a.py | 16 ++++- deepmd/pt/model/descriptor/dpa1.py | 35 ++++++++++- deepmd/pt/model/descriptor/dpa2.py | 32 +++++++++- deepmd/pt/model/descriptor/repformers.py | 2 +- deepmd/pt/model/descriptor/se_a.py | 44 ++++++++++--- deepmd/pt/model/descriptor/se_atten.py | 7 +-- deepmd/pt/model/model/dp_atomic_model.py | 69 +++++---------------- deepmd/pt/model/task/ener.py | 3 + source/tests/pt/test_permutation_denoise.py | 2 + source/tests/pt/test_rot_denoise.py | 2 + source/tests/pt/test_se_e2_a.py | 36 ++++++----- source/tests/pt/test_smooth_denoise.py | 2 + source/tests/pt/test_trans_denoise.py | 2 + 13 files changed, 164 insertions(+), 88 deletions(-) diff --git a/deepmd/model_format/se_e2_a.py b/deepmd/model_format/se_e2_a.py index fe516c8620..c766d9ea12 100644 --- a/deepmd/model_format/se_e2_a.py +++ b/deepmd/model_format/se_e2_a.py @@ -223,7 +223,16 @@ def call( Returns ------- descriptor - The descriptor. shape: nf x nloc x ng x axis_neuron + The descriptor. shape: nf x nloc x (ng x axis_neuron) + gr + The rotationally equivariant and permutationally invariant single particle + representation. shape: nf x nloc x ng x 3 + gg + The rotationally invariant pair-partical channel, this descriptor returns None + rr + The rotationally equivariant pair-partical channel, this descriptor returns None + ww + The smooth switch function. """ # nf x nloc x nnei x 4 rr, ww = self.env_mat.call(coord_ext, atype_ext, nlist, self.davg, self.dstd) @@ -238,15 +247,17 @@ def call( gg = self.cal_g(ss, tt) # nf x nloc x ng x 4 gr += np.einsum("flni,flnj->flij", gg, tr) + # nf x nloc x ng x 4 gr /= self.nnei gr1 = gr[:, :, : self.axis_neuron, :] # nf x nloc x ng x ng1 grrg = np.einsum("flid,fljd->flij", gr, gr1) # nf x nloc x (ng x ng1) grrg = grrg.reshape(nf, nloc, ng * self.axis_neuron) - return grrg + return grrg, gr[..., 1:], None, None, ww def serialize(self) -> dict: + """Serialize the descriptor to dict.""" return { "rcut": self.rcut, "rcut_smth": self.rcut_smth, @@ -271,6 +282,7 @@ def serialize(self) -> dict: @classmethod def deserialize(cls, data: dict) -> "DescrptSeA": + """Deserialize from dict.""" data = copy.deepcopy(data) variables = data.pop("@variables") embeddings = data.pop("embeddings") diff --git a/deepmd/pt/model/descriptor/dpa1.py b/deepmd/pt/model/descriptor/dpa1.py index dd34b815c9..8482f1cd9d 100644 --- a/deepmd/pt/model/descriptor/dpa1.py +++ b/deepmd/pt/model/descriptor/dpa1.py @@ -135,12 +135,42 @@ def forward( nlist: torch.Tensor, mapping: Optional[torch.Tensor] = None, ): + """Compute the descriptor. + + Parameters + ---------- + coord_ext + The extended coordinates of atoms. shape: nf x (nallx3) + atype_ext + The extended aotm types. shape: nf x nall + nlist + The neighbor list. shape: nf x nloc x nnei + mapping + The index mapping, not required by this descriptor. + + Returns + ------- + descriptor + The descriptor. shape: nf x nloc x (ng x axis_neuron) + gr + The rotationally equivariant and permutationally invariant single particle + representation. shape: nf x nloc x ng x 3 + g2 + The rotationally invariant pair-partical channel. + shape: nf x nloc x nnei x ng + h2 + The rotationally equivariant pair-partical channel. + shape: nf x nloc x nnei x 3 + sw + The smooth switch function. shape: nf x nloc x nnei + + """ del mapping nframes, nloc, nnei = nlist.shape nall = extended_coord.view(nframes, -1).shape[1] // 3 g1_ext = self.type_embedding(extended_atype) g1_inp = g1_ext[:, :nloc, :] - g1, env_mat, diff, rot_mat, sw = self.se_atten( + g1, g2, h2, rot_mat, sw = self.se_atten( nlist, extended_coord, extended_atype, @@ -149,4 +179,5 @@ def forward( ) if self.concat_output_tebd: g1 = torch.cat([g1, g1_inp], dim=-1) - return g1, env_mat, diff, rot_mat, sw + + return g1, rot_mat, g2, h2, sw diff --git a/deepmd/pt/model/descriptor/dpa2.py b/deepmd/pt/model/descriptor/dpa2.py index fbdbc91dd9..f70b88ac32 100644 --- a/deepmd/pt/model/descriptor/dpa2.py +++ b/deepmd/pt/model/descriptor/dpa2.py @@ -329,6 +329,36 @@ def forward( nlist: torch.Tensor, mapping: Optional[torch.Tensor] = None, ): + """Compute the descriptor. + + Parameters + ---------- + coord_ext + The extended coordinates of atoms. shape: nf x (nallx3) + atype_ext + The extended aotm types. shape: nf x nall + nlist + The neighbor list. shape: nf x nloc x nnei + mapping + The index mapping, mapps extended region index to local region. + + Returns + ------- + descriptor + The descriptor. shape: nf x nloc x (ng x axis_neuron) + gr + The rotationally equivariant and permutationally invariant single particle + representation. shape: nf x nloc x ng x 3 + gg + The rotationally invariant pair-partical channel. + shape: nf x nloc x nnei x ng + hh + The rotationally equivariant pair-partical channel. + shape: nf x nloc x nnei x 3 + sw + The smooth switch function. shape: nf x nloc x nnei + + """ nframes, nloc, nnei = nlist.shape nall = extended_coord.view(nframes, -1).shape[1] // 3 # nlists @@ -372,4 +402,4 @@ def forward( ) if self.concat_output_tebd: g1 = torch.cat([g1, g1_inp], dim=-1) - return g1, g2, h2, rot_mat, sw + return g1, rot_mat, g2, h2, sw diff --git a/deepmd/pt/model/descriptor/repformers.py b/deepmd/pt/model/descriptor/repformers.py index 26887b1b75..141b5dc745 100644 --- a/deepmd/pt/model/descriptor/repformers.py +++ b/deepmd/pt/model/descriptor/repformers.py @@ -256,7 +256,7 @@ def forward( # (nb x nloc) x ng2 x 3 rot_mat = torch.permute(h2g2, (0, 1, 3, 2)) - return g1, g2, h2, rot_mat.view(-1, self.dim_emb, 3), sw + return g1, g2, h2, rot_mat.view(-1, nloc, self.dim_emb, 3), sw def compute_input_stats(self, merged): """Update mean and stddev for descriptor elements.""" diff --git a/deepmd/pt/model/descriptor/se_a.py b/deepmd/pt/model/descriptor/se_a.py index 10aa66311e..53bd72f3e3 100644 --- a/deepmd/pt/model/descriptor/se_a.py +++ b/deepmd/pt/model/descriptor/se_a.py @@ -115,12 +115,40 @@ def get_data_process_key(cls, config): def forward( self, - extended_coord: torch.Tensor, - extended_atype: torch.Tensor, + coord_ext: torch.Tensor, + atype_ext: torch.Tensor, nlist: torch.Tensor, mapping: Optional[torch.Tensor] = None, ): - return self.sea.forward(nlist, extended_coord, extended_atype, None, mapping) + """Compute the descriptor. + + Parameters + ---------- + coord_ext + The extended coordinates of atoms. shape: nf x (nallx3) + atype_ext + The extended aotm types. shape: nf x nall + nlist + The neighbor list. shape: nf x nloc x nnei + mapping + The index mapping, not required by this descriptor. + + Returns + ------- + descriptor + The descriptor. shape: nf x nloc x (ng x axis_neuron) + gr + The rotationally equivariant and permutationally invariant single particle + representation. shape: nf x nloc x ng x 3 + gg + The rotationally invariant pair-partical channel, this descriptor returns None + rr + The rotationally equivariant pair-partical channel, this descriptor returns None + ww + The smooth switch function. + + """ + return self.sea.forward(nlist, coord_ext, atype_ext, None, mapping) def set_stat_mean_and_stddev( self, @@ -389,7 +417,7 @@ def forward( del extended_atype_embd, mapping nloc = nlist.shape[1] atype = extended_atype[:, :nloc] - dmatrix, diff, _ = prod_env_mat_se_a( + dmatrix, diff, sw = prod_env_mat_se_a( extended_coord, nlist, atype, @@ -438,12 +466,14 @@ def forward( result = torch.matmul( xyz_scatter_1, xyz_scatter_2 ) # shape is [nframes*nall, self.filter_neuron[-1], self.axis_neuron] + result = result.view(-1, nloc, self.filter_neuron[-1] * self.axis_neuron) + rot_mat = rot_mat.view([-1, nloc] + list(rot_mat.shape[1:])) # noqa:RUF005 return ( - result.view(-1, nloc, self.filter_neuron[-1] * self.axis_neuron), - None, - None, + result, + rot_mat, None, None, + sw, ) diff --git a/deepmd/pt/model/descriptor/se_atten.py b/deepmd/pt/model/descriptor/se_atten.py index 0c932f42f2..78cba59da7 100644 --- a/deepmd/pt/model/descriptor/se_atten.py +++ b/deepmd/pt/model/descriptor/se_atten.py @@ -281,9 +281,8 @@ def forward( self.rcut, self.rcut_smth, ) - dmatrix = dmatrix.view( - -1, self.ndescrpt - ) # shape is [nframes*nall, self.ndescrpt] + # [nfxnlocxnnei, self.ndescrpt] + dmatrix = dmatrix.view(-1, self.ndescrpt) nlist_mask = nlist != -1 nlist[nlist == -1] = 0 sw = torch.squeeze(sw, -1) @@ -328,7 +327,7 @@ def forward( return ( result.view(-1, nloc, self.filter_neuron[-1] * self.axis_neuron), ret.view(-1, nloc, self.nnei, self.filter_neuron[-1]), - diff, + dmatrix.view(-1, nloc, self.nnei, 4)[..., 1:], rot_mat.view(-1, self.filter_neuron[-1], 3), sw, ) diff --git a/deepmd/pt/model/model/dp_atomic_model.py b/deepmd/pt/model/model/dp_atomic_model.py index a0f9b25765..853eacb875 100644 --- a/deepmd/pt/model/model/dp_atomic_model.py +++ b/deepmd/pt/model/model/dp_atomic_model.py @@ -14,7 +14,6 @@ Descriptor, ) from deepmd.pt.model.task import ( - DenoiseNet, Fitting, ) @@ -93,40 +92,20 @@ def __init__( sampled=sampled, ) - # Fitting - if fitting_net: - fitting_net["type"] = fitting_net.get("type", "ener") - if self.descriptor_type not in ["se_e2_a"]: - fitting_net["ntypes"] = 1 - else: - fitting_net["ntypes"] = self.descriptor.get_ntype() - fitting_net["use_tebd"] = False - fitting_net["embedding_width"] = self.descriptor.dim_out - - self.grad_force = "direct" not in fitting_net["type"] - if not self.grad_force: - fitting_net["out_dim"] = self.descriptor.dim_emb - if "ener" in fitting_net["type"]: - fitting_net["return_energy"] = True - self.fitting_net = Fitting(**fitting_net) + fitting_net["type"] = fitting_net.get("type", "ener") + if self.descriptor_type not in ["se_e2_a"]: + fitting_net["ntypes"] = 1 else: - self.fitting_net = None - self.grad_force = False - if not self.split_nlist: - self.coord_denoise_net = DenoiseNet( - self.descriptor.dim_out, self.ntypes - 1, self.descriptor.dim_emb - ) - elif self.combination: - self.coord_denoise_net = DenoiseNet( - self.descriptor.dim_out, - self.ntypes - 1, - self.descriptor.dim_emb_list, - self.prefactor, - ) - else: - self.coord_denoise_net = DenoiseNet( - self.descriptor.dim_out, self.ntypes - 1, self.descriptor.dim_emb - ) + fitting_net["ntypes"] = self.descriptor.get_ntype() + fitting_net["use_tebd"] = False + fitting_net["embedding_width"] = self.descriptor.dim_out + + self.grad_force = "direct" not in fitting_net["type"] + if not self.grad_force: + fitting_net["out_dim"] = self.descriptor.dim_emb + if "ener" in fitting_net["type"]: + fitting_net["return_energy"] = True + self.fitting_net = Fitting(**fitting_net) def get_fitting_output_def(self) -> FittingOutputDef: """Get the output def of the fitting net.""" @@ -178,7 +157,7 @@ def forward_atomic( atype = extended_atype[:, :nloc] if self.do_grad(): extended_coord.requires_grad_(True) - descriptor, env_mat, diff, rot_mat, sw = self.descriptor( + descriptor, rot_mat, g2, h2, sw = self.descriptor( extended_coord, extended_atype, nlist, @@ -186,23 +165,5 @@ def forward_atomic( ) assert descriptor is not None # energy, force - if self.fitting_net is not None: - fit_ret = self.fitting_net( - descriptor, atype, atype_tebd=None, rot_mat=rot_mat - ) - # denoise - else: - nlist_list = [nlist] - if not self.split_nlist: - nnei_mask = nlist != -1 - elif self.combination: - nnei_mask = [] - for item in nlist_list: - nnei_mask_item = item != -1 - nnei_mask.append(nnei_mask_item) - else: - env_mat = env_mat[-1] - diff = diff[-1] - nnei_mask = nlist_list[-1] != -1 - fit_ret = self.coord_denoise_net(env_mat, diff, nnei_mask, descriptor, sw) + fit_ret = self.fitting_net(descriptor, atype, atype_tebd=None, rot_mat=rot_mat) return fit_ret diff --git a/deepmd/pt/model/task/ener.py b/deepmd/pt/model/task/ener.py index 7ddcbd5c54..03043e2fcb 100644 --- a/deepmd/pt/model/task/ener.py +++ b/deepmd/pt/model/task/ener.py @@ -207,8 +207,11 @@ def forward( inputs ) # Shape is [nframes, nloc, m1] assert list(vec_out.size()) == [nframes, nloc, self.out_dim] + # (nf x nloc) x 1 x od vec_out = vec_out.view(-1, 1, self.out_dim) assert rot_mat is not None + # (nf x nloc) x od x 3 + rot_mat = rot_mat.view(-1, self.out_dim, 3) vec_out = ( torch.bmm(vec_out, rot_mat).squeeze(-2).view(nframes, nloc, 3) ) # Shape is [nframes, nloc, 3] diff --git a/source/tests/pt/test_permutation_denoise.py b/source/tests/pt/test_permutation_denoise.py index 47bd0360f2..6dd61ab7e4 100644 --- a/source/tests/pt/test_permutation_denoise.py +++ b/source/tests/pt/test_permutation_denoise.py @@ -66,6 +66,7 @@ def test( ) +@unittest.skip("support of the denoise is temporally disabled") class TestDenoiseModelDPA1(unittest.TestCase, PermutationDenoiseTest): def setUp(self): model_params = copy.deepcopy(model_dpa1) @@ -74,6 +75,7 @@ def setUp(self): self.model = get_model(model_params, sampled).to(env.DEVICE) +@unittest.skip("support of the denoise is temporally disabled") class TestDenoiseModelDPA2(unittest.TestCase, PermutationDenoiseTest): def setUp(self): model_params_sample = copy.deepcopy(model_dpa2) diff --git a/source/tests/pt/test_rot_denoise.py b/source/tests/pt/test_rot_denoise.py index cab8de7bec..2cbfd8fd38 100644 --- a/source/tests/pt/test_rot_denoise.py +++ b/source/tests/pt/test_rot_denoise.py @@ -97,6 +97,7 @@ def test( ) +@unittest.skip("support of the denoise is temporally disabled") class TestDenoiseModelDPA1(unittest.TestCase, RotDenoiseTest): def setUp(self): model_params = copy.deepcopy(model_dpa1) @@ -105,6 +106,7 @@ def setUp(self): self.model = get_model(model_params, sampled).to(env.DEVICE) +@unittest.skip("support of the denoise is temporally disabled") class TestDenoiseModelDPA2(unittest.TestCase, RotDenoiseTest): def setUp(self): model_params_sample = copy.deepcopy(model_dpa2) diff --git a/source/tests/pt/test_se_e2_a.py b/source/tests/pt/test_se_e2_a.py index 96a17c2bad..c0a106cb16 100644 --- a/source/tests/pt/test_se_e2_a.py +++ b/source/tests/pt/test_se_e2_a.py @@ -102,7 +102,7 @@ def test_consistency( ) # serialization dd1 = DescrptSeA.deserialize(dd0.serialize()) - rd1, _, _, _, _ = dd1( + rd1, gr1, _, _, sw1 = dd1( torch.tensor(self.coord_ext, dtype=dtype, device=env.DEVICE), torch.tensor(self.atype_ext, dtype=int, device=env.DEVICE), torch.tensor(self.nlist, dtype=int, device=env.DEVICE), @@ -116,18 +116,19 @@ def test_consistency( ) # dp impl dd2 = DPDescrptSeA.deserialize(dd0.serialize()) - rd2 = dd2.call( + rd2, gr2, _, _, sw2 = dd2.call( self.coord_ext, self.atype_ext, self.nlist, ) - np.testing.assert_allclose( - rd0.detach().cpu().numpy(), - rd2, - rtol=rtol, - atol=atol, - err_msg=err_msg, - ) + for aa, bb in zip([rd1, gr1, sw1], [rd2, gr2, sw2]): + np.testing.assert_allclose( + aa.detach().cpu().numpy(), + bb, + rtol=rtol, + atol=atol, + err_msg=err_msg, + ) # old impl if idt is False and prec == "float64": dd3 = DescrptSeA( @@ -154,18 +155,19 @@ def test_consistency( dd3_state_dict[i] = dd3_state_dict[i].unsqueeze(0) dd3.sea.load_state_dict(dd3_state_dict) - rd3, _, _, _, _ = dd3( + rd3, gr3, _, _, sw3 = dd3( torch.tensor(self.coord_ext, dtype=dtype, device=env.DEVICE), torch.tensor(self.atype_ext, dtype=int, device=env.DEVICE), torch.tensor(self.nlist, dtype=int, device=env.DEVICE), ) - np.testing.assert_allclose( - rd0.detach().cpu().numpy(), - rd3.detach().cpu().numpy(), - rtol=rtol, - atol=atol, - err_msg=err_msg, - ) + for aa, bb in zip([rd1, gr1, sw1], [rd3, gr3, sw3]): + np.testing.assert_allclose( + aa.detach().cpu().numpy(), + bb.detach().cpu().numpy(), + rtol=rtol, + atol=atol, + err_msg=err_msg, + ) def test_jit( self, diff --git a/source/tests/pt/test_smooth_denoise.py b/source/tests/pt/test_smooth_denoise.py index a66e5df957..de89f8dccc 100644 --- a/source/tests/pt/test_smooth_denoise.py +++ b/source/tests/pt/test_smooth_denoise.py @@ -96,6 +96,7 @@ def compare(ret0, ret1): compare(ret0, ret3) +@unittest.skip("support of the denoise is temporally disabled") class TestDenoiseModelDPA2(unittest.TestCase, SmoothDenoiseTest): def setUp(self): model_params_sample = copy.deepcopy(model_dpa2) @@ -116,6 +117,7 @@ def setUp(self): self.aprec = 1e-5 +@unittest.skip("support of the denoise is temporally disabled") class TestDenoiseModelDPA2_1(unittest.TestCase, SmoothDenoiseTest): def setUp(self): model_params_sample = copy.deepcopy(model_dpa2) diff --git a/source/tests/pt/test_trans_denoise.py b/source/tests/pt/test_trans_denoise.py index 360633278c..88b926a3ae 100644 --- a/source/tests/pt/test_trans_denoise.py +++ b/source/tests/pt/test_trans_denoise.py @@ -56,6 +56,7 @@ def test( torch.testing.assert_close(ret0["logits"], ret1["logits"], rtol=prec, atol=prec) +@unittest.skip("support of the denoise is temporally disabled") class TestDenoiseModelDPA1(unittest.TestCase, TransDenoiseTest): def setUp(self): model_params = copy.deepcopy(model_dpa1) @@ -64,6 +65,7 @@ def setUp(self): self.model = get_model(model_params, sampled).to(env.DEVICE) +@unittest.skip("support of the denoise is temporally disabled") class TestDenoiseModelDPA2(unittest.TestCase, TransDenoiseTest): def setUp(self): model_params_sample = copy.deepcopy(model_dpa2) From e76cc661f0ba9e83cc707c68d4d921062767ee01 Mon Sep 17 00:00:00 2001 From: Han Wang Date: Sat, 27 Jan 2024 23:02:04 +0800 Subject: [PATCH 02/14] same output name for all descriptors --- deepmd/pt/model/descriptor/dpa1.py | 4 ++-- deepmd/pt/model/descriptor/dpa2.py | 8 ++++---- deepmd/pt/model/descriptor/se_a.py | 10 +++++----- 3 files changed, 11 insertions(+), 11 deletions(-) diff --git a/deepmd/pt/model/descriptor/dpa1.py b/deepmd/pt/model/descriptor/dpa1.py index 8482f1cd9d..23f521b6d8 100644 --- a/deepmd/pt/model/descriptor/dpa1.py +++ b/deepmd/pt/model/descriptor/dpa1.py @@ -156,10 +156,10 @@ def forward( The rotationally equivariant and permutationally invariant single particle representation. shape: nf x nloc x ng x 3 g2 - The rotationally invariant pair-partical channel. + The rotationally invariant pair-partical representation. shape: nf x nloc x nnei x ng h2 - The rotationally equivariant pair-partical channel. + The rotationally equivariant pair-partical representation. shape: nf x nloc x nnei x 3 sw The smooth switch function. shape: nf x nloc x nnei diff --git a/deepmd/pt/model/descriptor/dpa2.py b/deepmd/pt/model/descriptor/dpa2.py index f70b88ac32..409b999262 100644 --- a/deepmd/pt/model/descriptor/dpa2.py +++ b/deepmd/pt/model/descriptor/dpa2.py @@ -349,11 +349,11 @@ def forward( gr The rotationally equivariant and permutationally invariant single particle representation. shape: nf x nloc x ng x 3 - gg - The rotationally invariant pair-partical channel. + g2 + The rotationally invariant pair-partical representation. shape: nf x nloc x nnei x ng - hh - The rotationally equivariant pair-partical channel. + h2 + The rotationally equivariant pair-partical representation. shape: nf x nloc x nnei x 3 sw The smooth switch function. shape: nf x nloc x nnei diff --git a/deepmd/pt/model/descriptor/se_a.py b/deepmd/pt/model/descriptor/se_a.py index 53bd72f3e3..40b30e5c4b 100644 --- a/deepmd/pt/model/descriptor/se_a.py +++ b/deepmd/pt/model/descriptor/se_a.py @@ -140,11 +140,11 @@ def forward( gr The rotationally equivariant and permutationally invariant single particle representation. shape: nf x nloc x ng x 3 - gg - The rotationally invariant pair-partical channel, this descriptor returns None - rr - The rotationally equivariant pair-partical channel, this descriptor returns None - ww + g2 + The rotationally invariant pair-partical representation, this descriptor returns None + h2 + The rotationally equivariant pair-partical representation, this descriptor returns None + sw The smooth switch function. """ From eff9f697dd07a508adf607dbbae25e5efba114c6 Mon Sep 17 00:00:00 2001 From: Han Wang Date: Sat, 27 Jan 2024 23:04:01 +0800 Subject: [PATCH 03/14] update doc --- deepmd/pt/model/descriptor/se_a.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/deepmd/pt/model/descriptor/se_a.py b/deepmd/pt/model/descriptor/se_a.py index 40b30e5c4b..3f42736dca 100644 --- a/deepmd/pt/model/descriptor/se_a.py +++ b/deepmd/pt/model/descriptor/se_a.py @@ -141,9 +141,11 @@ def forward( The rotationally equivariant and permutationally invariant single particle representation. shape: nf x nloc x ng x 3 g2 - The rotationally invariant pair-partical representation, this descriptor returns None + The rotationally invariant pair-partical representation. + this descriptor returns None h2 - The rotationally equivariant pair-partical representation, this descriptor returns None + The rotationally equivariant pair-partical representation. + this descriptor returns None sw The smooth switch function. From 4cc9f708ad13964657e1f42d8c923bafc54383c9 Mon Sep 17 00:00:00 2001 From: Han Wang Date: Sat, 27 Jan 2024 23:04:16 +0800 Subject: [PATCH 04/14] update doc --- deepmd/model_format/se_e2_a.py | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/deepmd/model_format/se_e2_a.py b/deepmd/model_format/se_e2_a.py index c766d9ea12..28751cad8d 100644 --- a/deepmd/model_format/se_e2_a.py +++ b/deepmd/model_format/se_e2_a.py @@ -227,11 +227,13 @@ def call( gr The rotationally equivariant and permutationally invariant single particle representation. shape: nf x nloc x ng x 3 - gg - The rotationally invariant pair-partical channel, this descriptor returns None - rr - The rotationally equivariant pair-partical channel, this descriptor returns None - ww + g2 + The rotationally invariant pair-partical representation. + this descriptor returns None + h2 + The rotationally equivariant pair-partical representation. + this descriptor returns None + sw The smooth switch function. """ # nf x nloc x nnei x 4 From 480eab9011676d76e04af2a23b6eba2a552c0895 Mon Sep 17 00:00:00 2001 From: Han Wang Date: Sat, 27 Jan 2024 23:25:39 +0800 Subject: [PATCH 05/14] fix ut --- source/tests/test_model_format_utils.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/source/tests/test_model_format_utils.py b/source/tests/test_model_format_utils.py index 22393515ec..da76c53ed9 100644 --- a/source/tests/test_model_format_utils.py +++ b/source/tests/test_model_format_utils.py @@ -367,4 +367,5 @@ def test_self_consistency( em1 = DescrptSeA.deserialize(em0.serialize()) mm0 = em0.call(self.coord_ext, self.atype_ext, self.nlist) mm1 = em1.call(self.coord_ext, self.atype_ext, self.nlist) - np.testing.assert_allclose(mm0, mm1) + for ii in [0, 1, 4]: + np.testing.assert_allclose(mm0[ii], mm1[ii]) From b243bfc663676d65d905713a08c8b0a6c99edb39 Mon Sep 17 00:00:00 2001 From: Han Wang Date: Mon, 29 Jan 2024 15:49:48 +0800 Subject: [PATCH 06/14] add testing cases --- source/tests/test_model_format_utils.py | 102 ++++++++++++++++++++++++ 1 file changed, 102 insertions(+) diff --git a/source/tests/test_model_format_utils.py b/source/tests/test_model_format_utils.py index da76c53ed9..57564df1ae 100644 --- a/source/tests/test_model_format_utils.py +++ b/source/tests/test_model_format_utils.py @@ -13,6 +13,7 @@ EmbeddingNet, EnvMat, FittingNet, + InvarFitting, NativeLayer, NativeNet, NetworkCollection, @@ -369,3 +370,104 @@ def test_self_consistency( mm1 = em1.call(self.coord_ext, self.atype_ext, self.nlist) for ii in [0, 1, 4]: np.testing.assert_allclose(mm0[ii], mm1[ii]) + + +class TestInvarFitting(unittest.TestCase, TestCaseSingleFrameWithNlist): + def setUp(self): + TestCaseSingleFrameWithNlist.setUp(self) + + def test_self_consistency( + self, + ): + rng = np.random.default_rng() + nf, nloc, nnei = self.nlist.shape + ds = DescrptSeA(self.rcut, self.rcut_smth, self.sel) + dd = ds.call(self.coord_ext, self.atype_ext, self.nlist) + atype = self.atype_ext[:, :nloc] + + for ( + distinguish_types, + od, + nfp, + nap, + ) in itertools.product( + [True, False], + [1, 2], + [0, 3], + [0, 4], + ): + ifn0 = InvarFitting( + "energy", + self.nt, + ds.dim_out, + od, + numb_fparam=nfp, + numb_aparam=nap, + distinguish_types=distinguish_types, + ) + ifn1 = InvarFitting.deserialize(ifn0.serialize()) + if nfp > 0: + ifp = rng.normal(size=(self.nf, nfp)) + else: + ifp = None + if nap > 0: + iap = rng.normal(size=(self.nf, self.nloc, nap)) + else: + iap = None + ret0 = ifn0(dd[0], atype, fparam=ifp, aparam=iap) + ret1 = ifn1(dd[0], atype, fparam=ifp, aparam=iap) + np.testing.assert_allclose(ret0["energy"], ret1["energy"]) + + def test_self_exception( + self, + ): + rng = np.random.default_rng() + nf, nloc, nnei = self.nlist.shape + ds = DescrptSeA(self.rcut, self.rcut_smth, self.sel) + dd = ds.call(self.coord_ext, self.atype_ext, self.nlist) + atype = self.atype_ext[:, :nloc] + + for ( + distinguish_types, + od, + nfp, + nap, + ) in itertools.product( + [True, False], + [1, 2], + [0, 3], + [0, 4], + ): + ifn0 = InvarFitting( + "energy", + self.nt, + ds.dim_out, + od, + numb_fparam=nfp, + numb_aparam=nap, + distinguish_types=distinguish_types, + ) + + if nfp > 0: + ifp = rng.normal(size=(self.nf, nfp)) + else: + ifp = None + if nap > 0: + iap = rng.normal(size=(self.nf, self.nloc, nap)) + else: + iap = None + with self.assertRaises(ValueError) as context: + ret0 = ifn0(dd[0][:, :, :-2], atype, fparam=ifp, aparam=iap) + self.assertIn("input descriptor", context.exception) + + if nfp > 0: + ifp = rng.normal(size=(self.nf, nfp - 1)) + with self.assertRaises(ValueError) as context: + ret0 = ifn0(dd[0], atype, fparam=ifp, aparam=iap) + self.assertIn("input fparam", context.exception) + + if nap > 0: + iap = rng.normal(size=(self.nf, self.nloc, nap - 1)) + with self.assertRaises(ValueError) as context: + ret0 = ifn0(dd[0], atype, fparam=ifp, aparam=iap) + self.assertIn("input aparam", context.exception) From c49d8abf6c014f7b81b8dfcb41482a856e43e847 Mon Sep 17 00:00:00 2001 From: Han Wang Date: Mon, 29 Jan 2024 16:16:36 +0800 Subject: [PATCH 07/14] fix bug --- deepmd/model_format/se_e2_a.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/deepmd/model_format/se_e2_a.py b/deepmd/model_format/se_e2_a.py index 28751cad8d..f179b10ac3 100644 --- a/deepmd/model_format/se_e2_a.py +++ b/deepmd/model_format/se_e2_a.py @@ -171,9 +171,8 @@ def __init__( ) self.env_mat = EnvMat(self.rcut, self.rcut_smth) self.nnei = np.sum(self.sel) - self.nneix4 = self.nnei * 4 - self.davg = np.zeros([self.ntypes, self.nneix4]) - self.dstd = np.ones([self.ntypes, self.nneix4]) + self.davg = np.zeros([self.ntypes, self.nnei, 4]) + self.dstd = np.ones([self.ntypes, self.nnei, 4]) self.orig_sel = self.sel def __setitem__(self, key, value): @@ -192,6 +191,11 @@ def __getitem__(self, key): else: raise KeyError(key) + @property + def dim_out(self): + """Returns the output dimension of this descriptor.""" + return self.neuron[-1] * self.axis_neuron + def cal_g( self, ss, From b51a3c32e1ec733ad907ab13ca97c145dc981f51 Mon Sep 17 00:00:00 2001 From: Han Wang Date: Mon, 29 Jan 2024 22:11:36 +0800 Subject: [PATCH 08/14] add missing file for the model format --- deepmd/model_format/fitting.py | 356 +++++++++++++++++++++++++++++++++ 1 file changed, 356 insertions(+) create mode 100644 deepmd/model_format/fitting.py diff --git a/deepmd/model_format/fitting.py b/deepmd/model_format/fitting.py new file mode 100644 index 0000000000..b3195cd26e --- /dev/null +++ b/deepmd/model_format/fitting.py @@ -0,0 +1,356 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import copy +from typing import ( + Any, + List, + Optional, +) + +import numpy as np + +from deepmd.model_format import ( + FittingOutputDef, + OutputVariableDef, + fitting_check_output, +) + +from .common import ( + DEFAULT_PRECISION, + NativeOP, +) +from .network import ( + FittingNet, + NetworkCollection, +) + + +@fitting_check_output +class InvarFitting(NativeOP): + r"""Fitting the energy (or a porperty of `dim_out`) of the system. The force and the virial can also be trained. + + Lets take the energy fitting task as an example. + The potential energy :math:`E` is a fitting network function of the descriptor :math:`\mathcal{D}`: + + .. math:: + E(\mathcal{D}) = \mathcal{L}^{(n)} \circ \mathcal{L}^{(n-1)} + \circ \cdots \circ \mathcal{L}^{(1)} \circ \mathcal{L}^{(0)} + + The first :math:`n` hidden layers :math:`\mathcal{L}^{(0)}, \cdots, \mathcal{L}^{(n-1)}` are given by + + .. math:: + \mathbf{y}=\mathcal{L}(\mathbf{x};\mathbf{w},\mathbf{b})= + \boldsymbol{\phi}(\mathbf{x}^T\mathbf{w}+\mathbf{b}) + + where :math:`\mathbf{x} \in \mathbb{R}^{N_1}` is the input vector and :math:`\mathbf{y} \in \mathbb{R}^{N_2}` + is the output vector. :math:`\mathbf{w} \in \mathbb{R}^{N_1 \times N_2}` and + :math:`\mathbf{b} \in \mathbb{R}^{N_2}` are weights and biases, respectively, + both of which are trainable if `trainable[i]` is `True`. :math:`\boldsymbol{\phi}` + is the activation function. + + The output layer :math:`\mathcal{L}^{(n)}` is given by + + .. math:: + \mathbf{y}=\mathcal{L}^{(n)}(\mathbf{x};\mathbf{w},\mathbf{b})= + \mathbf{x}^T\mathbf{w}+\mathbf{b} + + where :math:`\mathbf{x} \in \mathbb{R}^{N_{n-1}}` is the input vector and :math:`\mathbf{y} \in \mathbb{R}` + is the output scalar. :math:`\mathbf{w} \in \mathbb{R}^{N_{n-1}}` and + :math:`\mathbf{b} \in \mathbb{R}` are weights and bias, respectively, + both of which are trainable if `trainable[n]` is `True`. + + Parameters + ---------- + var_name + The name of the output variable. + ntypes + The number of atom types. + dim_descrpt + The dimension of the input descriptor. + dim_out + The dimension of the output fit property. + neuron + Number of neurons :math:`N` in each hidden layer of the fitting net + resnet_dt + Time-step `dt` in the resnet construction: + :math:`y = x + dt * \phi (Wx + b)` + numb_fparam + Number of frame parameter + numb_aparam + Number of atomic parameter + rcond + The condition number for the regression of atomic energy. + tot_ener_zero + Force the total energy to zero. Useful for the charge fitting. + trainable + If the weights of fitting net are trainable. + Suppose that we have :math:`N_l` hidden layers in the fitting net, + this list is of length :math:`N_l + 1`, specifying if the hidden layers and the output layer are trainable. + atom_ener + Specifying atomic energy contribution in vacuum. The `set_davg_zero` key in the descrptor should be set. + activation_function + The activation function :math:`\boldsymbol{\phi}` in the embedding net. Supported options are |ACTIVATION_FN| + precision + The precision of the embedding net parameters. Supported options are |PRECISION| + layer_name : list[Optional[str]], optional + The name of the each layer. If two layers, either in the same fitting or different fittings, + have the same name, they will share the same neural network parameters. + use_aparam_as_mask: bool, optional + If True, the atomic parameters will be used as a mask that determines the atom is real/virtual. + And the aparam will not be used as the atomic parameters for embedding. + distinguish_types + Different atomic types uses different fitting net. + + """ + + def __init__( + self, + var_name: str, + ntypes: int, + dim_descrpt: int, + dim_out: int, + neuron: List[int] = [120, 120, 120], + resnet_dt: bool = True, + numb_fparam: int = 0, + numb_aparam: int = 0, + rcond: Optional[float] = None, + tot_ener_zero: bool = False, + trainable: Optional[List[bool]] = None, + atom_ener: Optional[List[float]] = None, + activation_function: str = "tanh", + precision: str = DEFAULT_PRECISION, + layer_name: Optional[List[Optional[str]]] = None, + use_aparam_as_mask: bool = False, + spin: Any = None, + distinguish_types: bool = False, + ): + # seed, uniform_seed are not included + if tot_ener_zero: + raise NotImplementedError("tot_ener_zero is not implemented") + if spin is not None: + raise NotImplementedError("spin is not implemented") + if use_aparam_as_mask: + raise NotImplementedError("use_aparam_as_mask is not implemented") + if use_aparam_as_mask: + raise NotImplementedError("use_aparam_as_mask is not implemented") + if layer_name is not None: + raise NotImplementedError("layer_name is not implemented") + if atom_ener is not None: + raise NotImplementedError("atom_ener is not implemented") + + self.var_name = var_name + self.ntypes = ntypes + self.dim_descrpt = dim_descrpt + self.dim_out = dim_out + self.neuron = neuron + self.resnet_dt = resnet_dt + self.numb_fparam = numb_fparam + self.numb_aparam = numb_aparam + self.rcond = rcond + self.tot_ener_zero = tot_ener_zero + self.trainable = trainable + self.atom_ener = atom_ener + self.activation_function = activation_function + self.precision = precision + self.layer_name = layer_name + self.use_aparam_as_mask = use_aparam_as_mask + self.spin = spin + self.distinguish_types = distinguish_types + if self.spin is not None: + raise NotImplementedError("spin is not supported") + + # init constants + self.bias_atom_e = np.zeros([self.ntypes, self.dim_out]) + if self.numb_fparam > 0: + self.fparam_avg = np.zeros(self.numb_fparam) + self.fparam_inv_std = np.ones(self.numb_fparam) + else: + self.fparam_avg, self.fparam_inv_std = None, None + if self.numb_aparam > 0: + self.aparam_avg = np.zeros(self.numb_aparam) + self.aparam_inv_std = np.ones(self.numb_aparam) + else: + self.aparam_avg, self.aparam_inv_std = None, None + # init networks + in_dim = self.dim_descrpt + self.numb_fparam + self.numb_aparam + out_dim = self.dim_out + self.nets = NetworkCollection( + 1 if self.distinguish_types else 0, + self.ntypes, + network_type="fitting_network", + networks=[ + FittingNet( + in_dim, + out_dim, + self.neuron, + self.activation_function, + self.resnet_dt, + self.precision, + bias_out=True, + ) + for ii in range(self.ntypes if self.distinguish_types else 1) + ], + ) + + def output_def(self): + return FittingOutputDef( + [ + OutputVariableDef( + self.var_name, [self.dim_out], reduciable=True, differentiable=True + ), + ] + ) + + def __setitem__(self, key, value): + if key in ("bias_atom_e"): + self.bias_atom_e = value + elif key in ("fparam_avg"): + self.fparam_avg = value + elif key in ("fparam_inv_std"): + self.fparam_inv_std = value + elif key in ("aparam_avg"): + self.aparam_avg = value + elif key in ("aparam_inv_std"): + self.aparam_inv_std = value + else: + raise KeyError(key) + + def __getitem__(self, key): + if key in ("bias_atom_e"): + return self.bias_atom_e + elif key in ("fparam_avg"): + return self.fparam_avg + elif key in ("fparam_inv_std"): + return self.fparam_inv_std + elif key in ("aparam_avg"): + return self.aparam_avg + elif key in ("aparam_inv_std"): + return self.aparam_inv_std + else: + raise KeyError(key) + + def serialize(self) -> dict: + """Serialize the fitting to dict.""" + return { + "var_name": self.var_name, + "ntypes": self.ntypes, + "dim_descrpt": self.dim_descrpt, + "dim_out": self.dim_out, + "neuron": self.neuron, + "resnet_dt": self.resnet_dt, + "numb_fparam": self.numb_fparam, + "numb_aparam": self.numb_aparam, + "rcond": self.rcond, + "activation_function": self.activation_function, + "precision": self.precision, + "distinguish_types": self.distinguish_types, + "nets": self.nets.serialize(), + "@variables": { + "bias_atom_e": self.bias_atom_e, + "fparam_avg": self.fparam_avg, + "fparam_inv_std": self.fparam_inv_std, + "aparam_avg": self.aparam_avg, + "aparam_inv_std": self.aparam_inv_std, + }, + # not supported + "tot_ener_zero": self.tot_ener_zero, + "trainable": self.trainable, + "atom_ener": self.atom_ener, + "layer_name": self.layer_name, + "use_aparam_as_mask": self.use_aparam_as_mask, + "spin": self.spin, + } + + @classmethod + def deserialize(cls, data: dict) -> "InvarFitting": + data = copy.deepcopy(data) + variables = data.pop("@variables") + nets = data.pop("nets") + obj = cls(**data) + for kk in variables.keys(): + obj[kk] = variables[kk] + obj.nets = NetworkCollection.deserialize(nets) + return obj + + def call( + self, + descriptor: np.array, + atype: np.array, + gr: Optional[np.array] = None, + g2: Optional[np.array] = None, + h2: Optional[np.array] = None, + fparam: Optional[np.array] = None, + aparam: Optional[np.array] = None, + ): + """Calculate the fitting. + + Parameters + ---------- + descriptor + input descriptor. shape: nf x nloc x nd + atype + the atom type. shape: nf x nloc + gr + The rotationally equivariant and permutationally invariant single particle + representation. shape: nf x nloc x ng x 3 + g2 + The rotationally invariant pair-partical representation. + shape: nf x nloc x nnei x ng + h2 + The rotationally equivariant pair-partical representation. + shape: nf x nloc x nnei x 3 + fparam + The frame parameter. shape: nf x nfp. nfp being `numb_fparam` + aparam + The atomic parameter. shape: nf x nloc x nap. nap being `numb_aparam` + + """ + nf, nloc, nd = descriptor.shape + # check input dim + if nd != self.dim_descrpt: + raise ValueError( + "get an input descriptor of dim {nd}," + "which is not consistent with {self.dim_descrpt}." + ) + xx = descriptor + # check fparam dim, concate to input descriptor + if self.numb_fparam > 0: + assert fparam is not None, "fparam should not be None" + if fparam.shape[-1] != self.numb_fparam: + raise ValueError( + "get an input fparam of dim {fparam.shape[-1]}, ", + "which is not consistent with {self.numb_fparam}.", + ) + fparam = (fparam - self.fparam_avg) * self.fparam_inv_std + fparam = np.tile(fparam.reshape([nf, 1, -1]), [1, nloc, 1]) + xx = np.concatenate( + [xx, fparam], + axis=-1, + ) + # check aparam dim, concate to input descriptor + if self.numb_aparam > 0: + assert aparam is not None, "aparam should not be None" + if aparam.shape[-1] != self.numb_aparam: + raise ValueError( + "get an input aparam of dim {aparam.shape[-1]}, ", + "which is not consistent with {self.numb_aparam}.", + ) + aparam = (aparam - self.aparam_avg) * self.aparam_inv_std + xx = np.concatenate( + [xx, aparam], + axis=-1, + ) + + # calcualte the prediction + if self.distinguish_types: + outs = np.zeros([nf, nloc, self.dim_out]) + for type_i in range(self.ntypes): + mask = np.tile( + (atype == type_i).reshape([nf, nloc, 1]), [1, 1, self.dim_out] + ) + atom_energy = self.nets[(type_i,)](xx) + atom_energy = atom_energy + self.bias_atom_e[type_i] + atom_energy = atom_energy * mask + outs = outs + atom_energy # Shape is [nframes, natoms[0], 1] + else: + outs = self.nets[()](xx) + self.bias_atom_e[atype] + return {self.var_name: outs} From 8afd47eb05f9a538dfce10fef0b92803d2ebd85c Mon Sep 17 00:00:00 2001 From: Han Wang Date: Mon, 29 Jan 2024 22:12:49 +0800 Subject: [PATCH 09/14] mlp should only use idt when a skip connection is available --- deepmd/model_format/network.py | 2 ++ deepmd/pt/model/network/mlp.py | 7 +++++-- 2 files changed, 7 insertions(+), 2 deletions(-) diff --git a/deepmd/model_format/network.py b/deepmd/model_format/network.py index a327d990c9..f2056c0b95 100644 --- a/deepmd/model_format/network.py +++ b/deepmd/model_format/network.py @@ -161,6 +161,8 @@ def __init__( ) -> None: prec = PRECISION_DICT[precision.lower()] self.precision = precision + # only use_timestep when skip connection is established. + use_timestep = use_timestep and (num_out == num_in or num_out == num_in * 2) rng = np.random.default_rng() self.w = rng.normal(size=(num_in, num_out)).astype(prec) self.b = rng.normal(size=(num_out,)).astype(prec) if bias else None diff --git a/deepmd/pt/model/network/mlp.py b/deepmd/pt/model/network/mlp.py index e3ac0e7bc2..d76abd82f9 100644 --- a/deepmd/pt/model/network/mlp.py +++ b/deepmd/pt/model/network/mlp.py @@ -56,7 +56,10 @@ def __init__( precision: str = DEFAULT_PRECISION, ): super().__init__() - self.use_timestep = use_timestep + # only use_timestep when skip connection is established. + self.use_timestep = use_timestep and ( + num_out == num_in or num_out == num_in * 2 + ) self.activate_name = activation_function self.activate = ActivationFn(self.activate_name) self.precision = precision @@ -207,7 +210,7 @@ class NetworkCollection(DPNetworkCollection, nn.Module): NETWORK_TYPE_MAP: ClassVar[Dict[str, type]] = { "network": MLP, "embedding_network": EmbeddingNet, - # "fitting_network": FittingNet, + "fitting_network": FittingNet, } def __init__(self, *args, **kwargs): From 19069f3f20d56fa46ef8a38d38d450befc9fb30c Mon Sep 17 00:00:00 2001 From: Han Wang Date: Mon, 29 Jan 2024 22:36:52 +0800 Subject: [PATCH 10/14] refactor the torch implementation of the fitting net --- deepmd/model_format/__init__.py | 4 + deepmd/model_format/fitting.py | 11 +- deepmd/pt/model/model/dp_atomic_model.py | 4 +- deepmd/pt/model/task/ener.py | 376 ++++++++++++++++++++--- deepmd/pt/model/task/fitting.py | 13 +- deepmd/pt/model/task/task.py | 18 +- deepmd/pt/utils/utils.py | 32 ++ source/tests/pt/test_ener_fitting.py | 162 ++++++++++ source/tests/pt/test_fitting_net.py | 24 +- source/tests/pt/test_model.py | 25 +- source/tests/pt/test_se_e2_a.py | 33 +- 11 files changed, 574 insertions(+), 128 deletions(-) create mode 100644 source/tests/pt/test_ener_fitting.py diff --git a/deepmd/model_format/__init__.py b/deepmd/model_format/__init__.py index 253bca3507..e15f73758e 100644 --- a/deepmd/model_format/__init__.py +++ b/deepmd/model_format/__init__.py @@ -7,6 +7,9 @@ from .env_mat import ( EnvMat, ) +from .fitting import ( + InvarFitting, +) from .network import ( EmbeddingNet, FittingNet, @@ -34,6 +37,7 @@ ) __all__ = [ + "InvarFitting", "DescrptSeA", "EnvMat", "make_multilayer_network", diff --git a/deepmd/model_format/fitting.py b/deepmd/model_format/fitting.py index b3195cd26e..8f79ae3491 100644 --- a/deepmd/model_format/fitting.py +++ b/deepmd/model_format/fitting.py @@ -8,12 +8,6 @@ import numpy as np -from deepmd.model_format import ( - FittingOutputDef, - OutputVariableDef, - fitting_check_output, -) - from .common import ( DEFAULT_PRECISION, NativeOP, @@ -22,6 +16,11 @@ FittingNet, NetworkCollection, ) +from .output_def import ( + FittingOutputDef, + OutputVariableDef, + fitting_check_output, +) @fitting_check_output diff --git a/deepmd/pt/model/model/dp_atomic_model.py b/deepmd/pt/model/model/dp_atomic_model.py index 853eacb875..245c0f3d3f 100644 --- a/deepmd/pt/model/model/dp_atomic_model.py +++ b/deepmd/pt/model/model/dp_atomic_model.py @@ -94,7 +94,7 @@ def __init__( fitting_net["type"] = fitting_net.get("type", "ener") if self.descriptor_type not in ["se_e2_a"]: - fitting_net["ntypes"] = 1 + fitting_net["ntypes"] = self.descriptor.get_ntype() else: fitting_net["ntypes"] = self.descriptor.get_ntype() fitting_net["use_tebd"] = False @@ -165,5 +165,5 @@ def forward_atomic( ) assert descriptor is not None # energy, force - fit_ret = self.fitting_net(descriptor, atype, atype_tebd=None, rot_mat=rot_mat) + fit_ret = self.fitting_net(descriptor, atype, gr=rot_mat) return fit_ret diff --git a/deepmd/pt/model/task/ener.py b/deepmd/pt/model/task/ener.py index 03043e2fcb..6b33491416 100644 --- a/deepmd/pt/model/task/ener.py +++ b/deepmd/pt/model/task/ener.py @@ -1,10 +1,13 @@ # SPDX-License-Identifier: LGPL-3.0-or-later +import copy import logging from typing import ( + List, Optional, Tuple, ) +import numpy as np import torch from deepmd.model_format import ( @@ -12,6 +15,10 @@ OutputVariableDef, fitting_check_output, ) +from deepmd.pt.model.network.mlp import ( + FittingNet, + NetworkCollection, +) from deepmd.pt.model.network.network import ( ResidualDeep, ) @@ -21,19 +28,35 @@ from deepmd.pt.utils import ( env, ) +from deepmd.pt.utils.env import ( + DEFAULT_PRECISION, + PRECISION_DICT, +) +from deepmd.pt.utils.utils import ( + to_numpy_array, + to_torch_tensor, +) + +dtype = env.GLOBAL_PT_FLOAT_PRECISION +device = env.DEVICE -@Fitting.register("ener") @fitting_check_output -class EnergyFittingNet(Fitting): +class InvarFitting(Fitting): def __init__( self, - ntypes, - embedding_width, - neuron, - bias_atom_e, - resnet_dt=True, - use_tebd=True, + var_name: str, + ntypes: int, + dim_descrpt: int, + dim_out: int, + neuron: List[int] = [128, 128, 128], + bias_atom_e: Optional[torch.Tensor] = None, + resnet_dt: bool = True, + numb_fparam: int = 0, + numb_aparam: int = 0, + activation_function: str = "tanh", + precision: str = DEFAULT_PRECISION, + distinguish_types: bool = False, **kwargs, ): """Construct a fitting net for energy. @@ -46,67 +69,325 @@ def __init__( - resnet_dt: Using time-step in the ResNet construction. """ super().__init__() + self.var_name = var_name self.ntypes = ntypes - self.embedding_width = embedding_width - self.use_tebd = use_tebd - if not use_tebd: - assert self.ntypes == len(bias_atom_e), "Element count mismatches!" - bias_atom_e = torch.tensor(bias_atom_e) + self.dim_descrpt = dim_descrpt + self.dim_out = dim_out + self.neuron = neuron + self.distinguish_types = distinguish_types + self.use_tebd = not self.distinguish_types + self.resnet_dt = resnet_dt + self.numb_fparam = numb_fparam + self.numb_aparam = numb_aparam + self.activation_function = activation_function + self.precision = precision + self.prec = PRECISION_DICT[self.precision] + if bias_atom_e is None: + bias_atom_e = np.zeros([self.ntypes, self.dim_out]) + bias_atom_e = torch.tensor(bias_atom_e, dtype=self.prec, device=device) + bias_atom_e = bias_atom_e.view([self.ntypes, self.dim_out]) + if not self.use_tebd: + assert self.ntypes == bias_atom_e.shape[0], "Element count mismatches!" self.register_buffer("bias_atom_e", bias_atom_e) + # init constants + if self.numb_fparam > 0: + self.register_buffer( + "fparam_avg", + torch.zeros(self.numb_fparam, dtype=self.prec, device=device), + ) + self.register_buffer( + "fparam_inv_std", + torch.ones(self.numb_fparam, dtype=self.prec, device=device), + ) + else: + self.fparam_avg, self.fparam_inv_std = None, None + if self.numb_aparam > 0: + self.register_buffer( + "aparam_avg", + torch.zeros(self.numb_aparam, dtype=self.prec, device=device), + ) + self.register_buffer( + "aparam_inv_std", + torch.ones(self.numb_aparam, dtype=self.prec, device=device), + ) + else: + self.aparam_avg, self.aparam_inv_std = None, None - filter_layers = [] - for type_i in range(self.ntypes): - bias_type = 0.0 - one = ResidualDeep( - type_i, embedding_width, neuron, bias_type, resnet_dt=resnet_dt + in_dim = self.dim_descrpt + self.numb_fparam + self.numb_aparam + out_dim = 1 + + self.old_impl = kwargs.get("old_impl", False) + if self.old_impl: + filter_layers = [] + for type_i in range(self.ntypes): + bias_type = 0.0 + one = ResidualDeep( + type_i, + self.dim_descrpt, + self.neuron, + bias_type, + resnet_dt=self.resnet_dt, + ) + filter_layers.append(one) + self.filter_layers_old = torch.nn.ModuleList(filter_layers) + self.filter_layers = None + else: + self.filter_layers = NetworkCollection( + 1 if self.distinguish_types else 0, + self.ntypes, + network_type="fitting_network", + networks=[ + FittingNet( + in_dim, + out_dim, + self.neuron, + self.activation_function, + self.resnet_dt, + self.precision, + bias_out=True, + ) + for ii in range(self.ntypes if self.distinguish_types else 1) + ], ) - filter_layers.append(one) - self.filter_layers = torch.nn.ModuleList(filter_layers) + self.filter_layers_old = None + # very bad design... if "seed" in kwargs: logging.info("Set seed to %d in fitting net.", kwargs["seed"]) torch.manual_seed(kwargs["seed"]) - def output_def(self): + def output_def(self) -> FittingOutputDef: return FittingOutputDef( [ - OutputVariableDef("energy", [1], reduciable=True, differentiable=True), + OutputVariableDef( + self.var_name, [self.dim_out], reduciable=True, differentiable=True + ), ] ) + def __setitem__(self, key, value): + if key in ("bias_atom_e"): + # correct bias_atom_e shape. user may provide stupid shape + self.bias_atom_e = value + elif key in ("fparam_avg"): + self.fparam_avg = value + elif key in ("fparam_inv_std"): + self.fparam_inv_std = value + elif key in ("aparam_avg"): + self.aparam_avg = value + elif key in ("aparam_inv_std"): + self.aparam_inv_std = value + else: + raise KeyError(key) + + def __getitem__(self, key): + if key in ("bias_atom_e"): + return self.bias_atom_e + elif key in ("fparam_avg"): + return self.fparam_avg + elif key in ("fparam_inv_std"): + return self.fparam_inv_std + elif key in ("aparam_avg"): + return self.aparam_avg + elif key in ("aparam_inv_std"): + return self.aparam_inv_std + else: + raise KeyError(key) + + def serialize(self) -> dict: + """Serialize the fitting to dict.""" + return { + "var_name": self.var_name, + "ntypes": self.ntypes, + "dim_descrpt": self.dim_descrpt, + "dim_out": self.dim_out, + "neuron": self.neuron, + "resnet_dt": self.resnet_dt, + "numb_fparam": self.numb_fparam, + "numb_aparam": self.numb_aparam, + "activation_function": self.activation_function, + "precision": self.precision, + "distinguish_types": self.distinguish_types, + "nets": self.filter_layers.serialize(), + "@variables": { + "bias_atom_e": to_numpy_array(self.bias_atom_e), + "fparam_avg": to_numpy_array(self.fparam_avg), + "fparam_inv_std": to_numpy_array(self.fparam_inv_std), + "aparam_avg": to_numpy_array(self.aparam_avg), + "aparam_inv_std": to_numpy_array(self.aparam_inv_std), + }, + # "rcond": self.rcond , + # "tot_ener_zero": self.tot_ener_zero , + # "trainable": self.trainable , + # "atom_ener": self.atom_ener , + # "layer_name": self.layer_name , + # "use_aparam_as_mask": self.use_aparam_as_mask , + # "spin": self.spin , + ## NOTICE: not supported by far + "rcond": None, + "tot_ener_zero": False, + "trainable": True, + "atom_ener": None, + "layer_name": None, + "use_aparam_as_mask": False, + "spin": None, + } + + @classmethod + def deserialize(cls, data: dict) -> "InvarFitting": + data = copy.deepcopy(data) + variables = data.pop("@variables") + nets = data.pop("nets") + obj = cls(**data) + for kk in variables.keys(): + obj[kk] = to_torch_tensor(variables[kk]) + obj.filter_layers = NetworkCollection.deserialize(nets) + return obj + + def _extend_f_avg_std(self, xx: torch.Tensor, nb: int) -> torch.Tensor: + return torch.tile(xx.view([1, self.numb_fparam]), [nb, 1]) + + def _extend_a_avg_std(self, xx: torch.Tensor, nb: int, nloc: int) -> torch.Tensor: + return torch.tile(xx.view([1, 1, self.numb_aparam]), [nb, nloc, 1]) + def forward( self, - inputs: torch.Tensor, + descriptor: torch.Tensor, atype: torch.Tensor, - atype_tebd: Optional[torch.Tensor] = None, - rot_mat: Optional[torch.Tensor] = None, + gr: Optional[torch.Tensor] = None, + g2: Optional[torch.Tensor] = None, + h2: Optional[torch.Tensor] = None, + fparam: Optional[torch.Tensor] = None, + aparam: Optional[torch.Tensor] = None, ): """Based on embedding net output, alculate total energy. Args: - - inputs: Embedding matrix. Its shape is [nframes, natoms[0], self.embedding_width]. + - inputs: Embedding matrix. Its shape is [nframes, natoms[0], self.dim_descrpt]. - natoms: Tell atom count and element count. Its shape is [2+self.ntypes]. Returns ------- - `torch.Tensor`: Total energy with shape [nframes, natoms[0]]. """ + xx = descriptor + nf, nloc, nd = xx.shape + dtype = descriptor.dtype + device = env.DEVICE + # NOTICE in tests/pt/test_model.py + # it happens that the user directly access the data memeber self.bias_atom_e + # and set it to a wrong shape! + self.bias_atom_e = self.bias_atom_e.view([self.ntypes, self.dim_out]) + # check input dim + if nd != self.dim_descrpt: + raise ValueError( + "get an input descriptor of dim {nd}," + "which is not consistent with {self.dim_descrpt}." + ) + # check fparam dim, concate to input descriptor + if self.numb_fparam > 0: + assert fparam is not None, "fparam should not be None" + assert self.fparam_avg is not None + assert self.fparam_inv_std is not None + if fparam.shape[-1] != self.numb_fparam: + raise ValueError( + "get an input fparam of dim {fparam.shape[-1]}, ", + "which is not consistent with {self.numb_fparam}.", + ) + nb, _ = fparam.shape + t_fparam_avg = self._extend_f_avg_std(self.fparam_avg, nb) + t_fparam_inv_std = self._extend_f_avg_std(self.fparam_inv_std, nb) + fparam = (fparam - t_fparam_avg) * t_fparam_inv_std + fparam = torch.tile(fparam.reshape([nf, 1, -1]), [1, nloc, 1]) + xx = torch.cat( + [xx, fparam], + dim=-1, + ) + # check aparam dim, concate to input descriptor + if self.numb_aparam > 0: + assert aparam is not None, "aparam should not be None" + assert self.aparam_avg is not None + assert self.aparam_inv_std is not None + if aparam.shape[-1] != self.numb_aparam: + raise ValueError( + "get an input aparam of dim {aparam.shape[-1]}, ", + "which is not consistent with {self.numb_aparam}.", + ) + nb, nloc, _ = aparam.shape + t_aparam_avg = self._extend_a_avg_std(self.aparam_avg, nb, nloc) + t_aparam_inv_std = self._extend_a_avg_std(self.aparam_inv_std, nb, nloc) + aparam = (aparam - t_aparam_avg) * t_aparam_inv_std + xx = torch.cat( + [xx, aparam], + dim=-1, + ) + outs = torch.zeros_like(atype).unsqueeze(-1) # jit assertion - if self.use_tebd: - if atype_tebd is not None: - inputs = torch.concat([inputs, atype_tebd], dim=-1) - atom_energy = self.filter_layers[0](inputs) + self.bias_atom_e[ - atype - ].unsqueeze(-1) - outs = outs + atom_energy # Shape is [nframes, natoms[0], 1] + if self.old_impl: + outs = torch.zeros_like(atype).unsqueeze(-1) # jit assertion + assert self.filter_layers_old is not None + if self.use_tebd: + atom_energy = self.filter_layers_old[0](xx) + self.bias_atom_e[ + atype + ].unsqueeze(-1) + outs = outs + atom_energy # Shape is [nframes, natoms[0], 1] + else: + for type_i, filter_layer in enumerate(self.filter_layers_old): + mask = atype == type_i + atom_energy = filter_layer(xx) + atom_energy = atom_energy + self.bias_atom_e[type_i] + atom_energy = atom_energy * mask.unsqueeze(-1) + outs = outs + atom_energy # Shape is [nframes, natoms[0], 1] + return {"energy": outs.to(env.GLOBAL_PT_FLOAT_PRECISION)} else: - for type_i, filter_layer in enumerate(self.filter_layers): - mask = atype == type_i - atom_energy = filter_layer(inputs) - atom_energy = atom_energy + self.bias_atom_e[type_i] - atom_energy = atom_energy * mask.unsqueeze(-1) + if self.use_tebd: + atom_energy = ( + self.filter_layers.networks[0](xx) + self.bias_atom_e[atype] + ) outs = outs + atom_energy # Shape is [nframes, natoms[0], 1] - return {"energy": outs.to(env.GLOBAL_PT_FLOAT_PRECISION)} + else: + for type_i, ll in enumerate(self.filter_layers.networks): + mask = (atype == type_i).unsqueeze(-1) + mask = torch.tile(mask, (1, 1, self.dim_out)) + atom_energy = ll(xx) + atom_energy = atom_energy + self.bias_atom_e[type_i] + atom_energy = atom_energy * mask + outs = outs + atom_energy # Shape is [nframes, natoms[0], 1] + return {self.var_name: outs.to(env.GLOBAL_PT_FLOAT_PRECISION)} + + +@Fitting.register("ener") +@fitting_check_output +class EnergyFittingNet(InvarFitting): + def __init__( + self, + ntypes: int, + embedding_width: int, + neuron: List[int] = [128, 128, 128], + bias_atom_e: Optional[torch.Tensor] = None, + resnet_dt: bool = True, + numb_fparam: int = 0, + numb_aparam: int = 0, + activation_function: str = "tanh", + precision: str = DEFAULT_PRECISION, + use_tebd: bool = True, + **kwargs, + ): + super().__init__( + "energy", + ntypes, + embedding_width, + 1, + neuron=neuron, + bias_atom_e=bias_atom_e, + resnet_dt=resnet_dt, + numb_fparam=numb_fparam, + numb_aparam=numb_aparam, + activation_function=activation_function, + precision=precision, + use_tebd=use_tebd, + **kwargs, + ) @Fitting.register("direct_force") @@ -136,7 +417,7 @@ def __init__( """ super().__init__() self.ntypes = ntypes - self.embedding_width = embedding_width + self.dim_descrpt = embedding_width self.use_tebd = use_tebd self.out_dim = out_dim if not use_tebd: @@ -186,13 +467,12 @@ def forward( self, inputs: torch.Tensor, atype: torch.Tensor, - atype_tebd: Optional[torch.Tensor] = None, - rot_mat: Optional[torch.Tensor] = None, + gr: Optional[torch.Tensor] = None, ) -> Tuple[torch.Tensor, None]: """Based on embedding net output, alculate total energy. Args: - - inputs: Embedding matrix. Its shape is [nframes, natoms[0], self.embedding_width]. + - inputs: Embedding matrix. Its shape is [nframes, natoms[0], self.dim_descrpt]. - natoms: Tell atom count and element count. Its shape is [2+self.ntypes]. Returns @@ -201,19 +481,19 @@ def forward( """ nframes, nloc, _ = inputs.size() if self.use_tebd: - if atype_tebd is not None: - inputs = torch.concat([inputs, atype_tebd], dim=-1) + # if atype_tebd is not None: + # inputs = torch.concat([inputs, atype_tebd], dim=-1) vec_out = self.filter_layers_dipole[0]( inputs ) # Shape is [nframes, nloc, m1] assert list(vec_out.size()) == [nframes, nloc, self.out_dim] # (nf x nloc) x 1 x od vec_out = vec_out.view(-1, 1, self.out_dim) - assert rot_mat is not None + assert gr is not None # (nf x nloc) x od x 3 - rot_mat = rot_mat.view(-1, self.out_dim, 3) + gr = gr.view(-1, self.out_dim, 3) vec_out = ( - torch.bmm(vec_out, rot_mat).squeeze(-2).view(nframes, nloc, 3) + torch.bmm(vec_out, gr).squeeze(-2).view(nframes, nloc, 3) ) # Shape is [nframes, nloc, 3] else: vec_out = torch.zeros_like(atype).unsqueeze(-1) # jit assertion diff --git a/deepmd/pt/model/task/fitting.py b/deepmd/pt/model/task/fitting.py index 16e80f9c20..c6fb6b27e1 100644 --- a/deepmd/pt/model/task/fitting.py +++ b/deepmd/pt/model/task/fitting.py @@ -7,9 +7,6 @@ import numpy as np import torch -from deepmd.model_format import ( - FittingOutputDef, -) from deepmd.pt.model.task.task import ( TaskBaseMethod, ) @@ -61,17 +58,9 @@ def __new__(cls, *args, **kwargs): if fitting_type in Fitting.__plugins.plugins: cls = Fitting.__plugins.plugins[fitting_type] else: - raise RuntimeError("Unknown descriptor type: " + fitting_type) + raise RuntimeError("Unknown fitting type: " + fitting_type) return super().__new__(cls) - def output_def(self) -> FittingOutputDef: - """Definition for the task Output.""" - raise NotImplementedError - - def forward(self, **kwargs): - """Task Output.""" - raise NotImplementedError - def share_params(self, base_class, shared_level, resume=False): assert ( self.__class__ == base_class.__class__ diff --git a/deepmd/pt/model/task/task.py b/deepmd/pt/model/task/task.py index a9b2efeb9a..b2dc03e4bd 100644 --- a/deepmd/pt/model/task/task.py +++ b/deepmd/pt/model/task/task.py @@ -1,12 +1,18 @@ # SPDX-License-Identifier: LGPL-3.0-or-later +from abc import ( + ABC, + abstractmethod, +) + import torch +from deepmd.model_format import ( + FittingOutputDef, +) -class TaskBaseMethod(torch.nn.Module): - def __init__(self, **kwargs): - """Construct a basic head for different tasks.""" - super().__init__() - def forward(self, **kwargs): - """Task Output.""" +class TaskBaseMethod(torch.nn.Module, ABC): + @abstractmethod + def output_def(self) -> FittingOutputDef: + """Definition for the task Output.""" raise NotImplementedError diff --git a/deepmd/pt/utils/utils.py b/deepmd/pt/utils/utils.py index 780dbf7e62..516cbbdba6 100644 --- a/deepmd/pt/utils/utils.py +++ b/deepmd/pt/utils/utils.py @@ -4,9 +4,17 @@ Optional, ) +import numpy as np import torch import torch.nn.functional as F +from deepmd.model_format.common import PRECISION_DICT as NP_PRECISION_DICT + +from .env import ( + DEVICE, +) +from .env import PRECISION_DICT as PT_PRECISION_DICT + def get_activation_fn(activation: str) -> Callable: """Returns the activation function corresponding to `activation`.""" @@ -41,3 +49,27 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: return x else: raise RuntimeError(f"activation function {self.activation} not supported") + + +def to_numpy_array( + xx: torch.Tensor, +) -> np.ndarray: + if xx is not None: + prec = [key for key, value in PT_PRECISION_DICT.items() if value == xx.dtype] + if len(prec) == 0: + raise ValueError(f"unknown precision {xx.dtype}") + else: + prec = NP_PRECISION_DICT[prec[0]] + return xx.detach().cpu().numpy().astype(prec) if xx is not None else None + + +def to_torch_tensor( + xx: np.ndarray, +) -> torch.Tensor: + if xx is not None: + prec = [key for key, value in NP_PRECISION_DICT.items() if value == xx.dtype] + if len(prec) == 0: + raise ValueError(f"unknown precision {xx.dtype}") + else: + prec = PT_PRECISION_DICT[prec[0]] + return torch.tensor(xx, dtype=prec, device=DEVICE) if xx is not None else None diff --git a/source/tests/pt/test_ener_fitting.py b/source/tests/pt/test_ener_fitting.py new file mode 100644 index 0000000000..f0d6bdb932 --- /dev/null +++ b/source/tests/pt/test_ener_fitting.py @@ -0,0 +1,162 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import itertools +import unittest + +import numpy as np +import torch + +from deepmd.model_format import InvarFitting as DPInvarFitting +from deepmd.pt.model.descriptor.se_a import ( + DescrptSeA, +) +from deepmd.pt.model.task.ener import ( + EnergyFittingNet, + InvarFitting, +) +from deepmd.pt.utils import ( + env, +) +from deepmd.pt.utils.utils import ( + to_numpy_array, +) + +from .test_env_mat import ( + TestCaseSingleFrameWithNlist, +) + +dtype = env.GLOBAL_PT_FLOAT_PRECISION + + +class TestInvarFitting(unittest.TestCase, TestCaseSingleFrameWithNlist): + def setUp(self): + TestCaseSingleFrameWithNlist.setUp(self) + + def test_consistency( + self, + ): + rng = np.random.default_rng() + nf, nloc, nnei = self.nlist.shape + dd0 = DescrptSeA(self.rcut, self.rcut_smth, self.sel).to(env.DEVICE) + rd0, _, _, _, _ = dd0( + torch.tensor(self.coord_ext, dtype=dtype, device=env.DEVICE), + torch.tensor(self.atype_ext, dtype=int, device=env.DEVICE), + torch.tensor(self.nlist, dtype=int, device=env.DEVICE), + ) + atype = torch.tensor(self.atype_ext[:, :nloc], dtype=int, device=env.DEVICE) + + for od, distinguish_types, nfp, nap in itertools.product( + [1, 3], + [True, False], + [0, 3], + [0, 4], + ): + ft0 = InvarFitting( + "foo", + self.nt, + dd0.dim_out, + od, + numb_fparam=nfp, + numb_aparam=nap, + use_tebd=(not distinguish_types), + ).to(env.DEVICE) + ft1 = DPInvarFitting.deserialize(ft0.serialize()) + ft2 = InvarFitting.deserialize(ft0.serialize()) + + if nfp > 0: + ifp = torch.tensor( + rng.normal(size=(self.nf, nfp)), dtype=dtype, device=env.DEVICE + ) + else: + ifp = None + if nap > 0: + iap = torch.tensor( + rng.normal(size=(self.nf, self.nloc, nap)), + dtype=dtype, + device=env.DEVICE, + ) + else: + iap = None + + ret0 = ft0(rd0, atype, fparam=ifp, aparam=iap) + ret1 = ft1( + rd0.detach().cpu().numpy(), + atype.detach().cpu().numpy(), + fparam=to_numpy_array(ifp), + aparam=to_numpy_array(iap), + ) + ret2 = ft2(rd0, atype, fparam=ifp, aparam=iap) + np.testing.assert_allclose( + to_numpy_array(ret0["foo"]), + ret1["foo"], + ) + np.testing.assert_allclose( + to_numpy_array(ret0["foo"]), + to_numpy_array(ret2["foo"]), + ) + + def test_new_old( + self, + ): + rng = np.random.default_rng() + nf, nloc, nnei = self.nlist.shape + dd = DescrptSeA(self.rcut, self.rcut_smth, self.sel).to(env.DEVICE) + rd0, _, _, _, _ = dd( + torch.tensor(self.coord_ext, dtype=dtype, device=env.DEVICE), + torch.tensor(self.atype_ext, dtype=int, device=env.DEVICE), + torch.tensor(self.nlist, dtype=int, device=env.DEVICE), + ) + atype = torch.tensor(self.atype_ext[:, :nloc], dtype=int, device=env.DEVICE) + + od = 1 + for distinguish_types in itertools.product( + [True, False], + ): + ft0 = EnergyFittingNet( + self.nt, + dd.dim_out, + distinguish_types=distinguish_types, + ).to(env.DEVICE) + ft1 = EnergyFittingNet( + self.nt, + dd.dim_out, + distinguish_types=distinguish_types, + old_impl=True, + ).to(env.DEVICE) + dd0 = ft0.state_dict() + dd1 = ft1.state_dict() + for kk, vv in dd1.items(): + new_kk = kk + new_kk = new_kk.replace("filter_layers_old", "filter_layers.networks") + new_kk = new_kk.replace("deep_layers", "layers") + new_kk = new_kk.replace("final_layer", "layers.3") + dd1[kk] = dd0[new_kk] + if kk.split(".")[-1] in ["idt", "bias"]: + dd1[kk] = dd1[kk].unsqueeze(0) + dd1["bias_atom_e"] = dd0["bias_atom_e"] + ft1.load_state_dict(dd1) + ret0 = ft0(rd0, atype) + ret1 = ft1(rd0, atype) + np.testing.assert_allclose( + to_numpy_array(ret0["energy"]), + to_numpy_array(ret1["energy"]), + ) + + def test_jit( + self, + ): + for od, distinguish_types, nfp, nap in itertools.product( + [1, 3], + [True, False], + [0, 3], + [0, 4], + ): + ft0 = InvarFitting( + "foo", + self.nt, + 9, + od, + numb_fparam=nfp, + numb_aparam=nap, + use_tebd=(not distinguish_types), + ).to(env.DEVICE) + torch.jit.script(ft0) diff --git a/source/tests/pt/test_fitting_net.py b/source/tests/pt/test_fitting_net.py index 3feb4f4739..ed2c428de5 100644 --- a/source/tests/pt/test_fitting_net.py +++ b/source/tests/pt/test_fitting_net.py @@ -102,25 +102,25 @@ def test_consistency(self): my_fn = EnergyFittingNet( self.ntypes, self.embedding_width, - self.n_neuron, - self.dp_fn.bias_atom_e, - use_tebd=False, + neuron=self.n_neuron, + bias_atom_e=self.dp_fn.bias_atom_e, + distinguish_types=True, ) for name, param in my_fn.named_parameters(): - matched = re.match("filter_layers\.(\d).deep_layers\.(\d)\.([a-z]+)", name) + matched = re.match( + "filter_layers\.networks\.(\d).layers\.(\d)\.([a-z]+)", name + ) key = None if matched: + if int(matched.group(2)) == len(self.n_neuron): + layer_id = -1 + else: + layer_id = matched.group(2) key = gen_key( type_id=matched.group(1), - layer_id=matched.group(2), + layer_id=layer_id, w_or_b=matched.group(3), ) - else: - matched = re.match("filter_layers\.(\d).final_layer\.([a-z]+)", name) - if matched: - key = gen_key( - type_id=matched.group(1), layer_id=-1, w_or_b=matched.group(2) - ) assert key is not None var = values[key] with torch.no_grad(): @@ -132,7 +132,7 @@ def test_consistency(self): ret = my_fn(embedding, atype) my_energy = ret["energy"] my_energy = my_energy.detach() - self.assertTrue(np.allclose(dp_energy, my_energy.numpy().reshape([-1]))) + np.testing.assert_allclose(dp_energy, my_energy.numpy().reshape([-1])) if __name__ == "__main__": diff --git a/source/tests/pt/test_model.py b/source/tests/pt/test_model.py index 5bbbc9e352..c6595e6471 100644 --- a/source/tests/pt/test_model.py +++ b/source/tests/pt/test_model.py @@ -53,23 +53,24 @@ VariableState = collections.namedtuple("VariableState", ["value", "gradient"]) -def torch2tf(torch_name): +def torch2tf(torch_name, last_layer_id=None): fields = torch_name.split(".") offset = int(fields[2] == "networks") element_id = int(fields[2 + offset]) if fields[0] == "descriptor": layer_id = int(fields[4 + offset]) + 1 weight_type = fields[5 + offset] - return "filter_type_all/%s_%d_%d:0" % (weight_type, layer_id, element_id) - elif fields[3] == "deep_layers": - layer_id = int(fields[4]) - weight_type = fields[5] - return "layer_%d_type_%d/%s:0" % (layer_id, element_id, weight_type) - elif fields[3] == "final_layer": - weight_type = fields[4] - return "final_layer_type_%d/%s:0" % (element_id, weight_type) + ret = "filter_type_all/%s_%d_%d:0" % (weight_type, layer_id, element_id) + elif fields[0] == "fitting_net": + layer_id = int(fields[4 + offset]) + weight_type = fields[5 + offset] + if layer_id != last_layer_id: + ret = "layer_%d_type_%d/%s:0" % (layer_id, element_id, weight_type) + else: + ret = "final_layer_type_%d/%s:0" % (element_id, weight_type) else: raise RuntimeError("Unexpected parameter name: %s" % torch_name) + return ret class DpTrainer: @@ -290,7 +291,7 @@ def test_consistency(self): "neuron": self.filter_neuron, "axis_neuron": self.axis_neuron, }, - "fitting_net": {"neuron": self.n_neuron}, + "fitting_net": {"neuron": self.n_neuron, "distinguish_types": True}, "data_stat_nbatch": self.data_stat_nbatch, "type_map": self.type_map, }, @@ -323,7 +324,7 @@ def test_consistency(self): # Keep parameter value consistency between 2 implentations for name, param in my_model.named_parameters(): name = name.replace("sea.", "") - var_name = torch2tf(name) + var_name = torch2tf(name, last_layer_id=len(self.n_neuron)) var = vs_dict[var_name].value with torch.no_grad(): src = torch.from_numpy(var) @@ -404,7 +405,7 @@ def step(step_id): for name, param in my_model.named_parameters(): name = name.replace("sea.", "") - var_name = torch2tf(name) + var_name = torch2tf(name, last_layer_id=len(self.n_neuron)) var_grad = vs_dict[var_name].gradient param_grad = param.grad.cpu() var_grad = torch.tensor(var_grad) diff --git a/source/tests/pt/test_se_e2_a.py b/source/tests/pt/test_se_e2_a.py index c0a106cb16..0da80ea1ea 100644 --- a/source/tests/pt/test_se_e2_a.py +++ b/source/tests/pt/test_se_e2_a.py @@ -25,6 +25,9 @@ PRECISION_DICT, ) +from .test_env_mat import ( + TestCaseSingleFrameWithNlist, +) from .test_mlp import ( get_tols, ) @@ -32,36 +35,6 @@ dtype = env.GLOBAL_PT_FLOAT_PRECISION -class TestCaseSingleFrameWithNlist: - def setUp(self): - # nloc == 3, nall == 4 - self.nloc = 3 - self.nall = 4 - self.nf, self.nt = 1, 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]) - # sel = [5, 2] - self.sel = [5, 2] - self.nlist = np.array( - [ - [1, 3, -1, -1, -1, 2, -1], - [0, -1, -1, -1, -1, 2, -1], - [0, 1, -1, -1, -1, 0, -1], - ], - dtype=int, - ).reshape([1, self.nloc, sum(self.sel)]) - self.rcut = 0.4 - self.rcut_smth = 2.2 - - # to be merged with the tf test case @unittest.skipIf(not support_se_e2_a, "EnvMat not supported") class TestDescrptSeA(unittest.TestCase, TestCaseSingleFrameWithNlist): From 49a10a9e377f8d656027d87742aec46589ce46e5 Mon Sep 17 00:00:00 2001 From: Han Wang Date: Tue, 30 Jan 2024 08:50:07 +0800 Subject: [PATCH 11/14] fix duplicated lines. fix use_tebd that is deprecated. --- deepmd/pt/model/model/dp_atomic_model.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/deepmd/pt/model/model/dp_atomic_model.py b/deepmd/pt/model/model/dp_atomic_model.py index 245c0f3d3f..a222c8e6f6 100644 --- a/deepmd/pt/model/model/dp_atomic_model.py +++ b/deepmd/pt/model/model/dp_atomic_model.py @@ -93,11 +93,11 @@ def __init__( ) fitting_net["type"] = fitting_net.get("type", "ener") - if self.descriptor_type not in ["se_e2_a"]: - fitting_net["ntypes"] = self.descriptor.get_ntype() + fitting_net["ntypes"] = self.descriptor.get_ntype() + if self.descriptor_type in ["se_e2_a"]: + fitting_net["distinguish_types"] = True else: - fitting_net["ntypes"] = self.descriptor.get_ntype() - fitting_net["use_tebd"] = False + fitting_net["distinguish_types"] = False fitting_net["embedding_width"] = self.descriptor.dim_out self.grad_force = "direct" not in fitting_net["type"] From 6efab8c203c14281fca2cff9b23778ae1ae42d85 Mon Sep 17 00:00:00 2001 From: Han Wang Date: Tue, 30 Jan 2024 09:00:02 +0800 Subject: [PATCH 12/14] find element in list rather than tuple, which easily makes bugs --- deepmd/model_format/fitting.py | 20 +++++++++---------- deepmd/pt/model/task/ener.py | 20 +++++++++---------- .../tests/common/test_model_format_utils.py | 19 ++++++++++++++++++ source/tests/pt/test_ener_fitting.py | 19 ++++++++++++++++++ 4 files changed, 58 insertions(+), 20 deletions(-) diff --git a/deepmd/model_format/fitting.py b/deepmd/model_format/fitting.py index 8f79ae3491..904fb42b76 100644 --- a/deepmd/model_format/fitting.py +++ b/deepmd/model_format/fitting.py @@ -200,29 +200,29 @@ def output_def(self): ) def __setitem__(self, key, value): - if key in ("bias_atom_e"): + if key in ["bias_atom_e"]: self.bias_atom_e = value - elif key in ("fparam_avg"): + elif key in ["fparam_avg"]: self.fparam_avg = value - elif key in ("fparam_inv_std"): + elif key in ["fparam_inv_std"]: self.fparam_inv_std = value - elif key in ("aparam_avg"): + elif key in ["aparam_avg"]: self.aparam_avg = value - elif key in ("aparam_inv_std"): + elif key in ["aparam_inv_std"]: self.aparam_inv_std = value else: raise KeyError(key) def __getitem__(self, key): - if key in ("bias_atom_e"): + if key in ["bias_atom_e"]: return self.bias_atom_e - elif key in ("fparam_avg"): + elif key in ["fparam_avg"]: return self.fparam_avg - elif key in ("fparam_inv_std"): + elif key in ["fparam_inv_std"]: return self.fparam_inv_std - elif key in ("aparam_avg"): + elif key in ["aparam_avg"]: return self.aparam_avg - elif key in ("aparam_inv_std"): + elif key in ["aparam_inv_std"]: return self.aparam_inv_std else: raise KeyError(key) diff --git a/deepmd/pt/model/task/ener.py b/deepmd/pt/model/task/ener.py index 6b33491416..1d241469c2 100644 --- a/deepmd/pt/model/task/ener.py +++ b/deepmd/pt/model/task/ener.py @@ -166,30 +166,30 @@ def output_def(self) -> FittingOutputDef: ) def __setitem__(self, key, value): - if key in ("bias_atom_e"): + if key in ["bias_atom_e"]: # correct bias_atom_e shape. user may provide stupid shape self.bias_atom_e = value - elif key in ("fparam_avg"): + elif key in ["fparam_avg"]: self.fparam_avg = value - elif key in ("fparam_inv_std"): + elif key in ["fparam_inv_std"]: self.fparam_inv_std = value - elif key in ("aparam_avg"): + elif key in ["aparam_avg"]: self.aparam_avg = value - elif key in ("aparam_inv_std"): + elif key in ["aparam_inv_std"]: self.aparam_inv_std = value else: raise KeyError(key) def __getitem__(self, key): - if key in ("bias_atom_e"): + if key in ["bias_atom_e"]: return self.bias_atom_e - elif key in ("fparam_avg"): + elif key in ["fparam_avg"]: return self.fparam_avg - elif key in ("fparam_inv_std"): + elif key in ["fparam_inv_std"]: return self.fparam_inv_std - elif key in ("aparam_avg"): + elif key in ["aparam_avg"]: return self.aparam_avg - elif key in ("aparam_inv_std"): + elif key in ["aparam_inv_std"]: return self.aparam_inv_std else: raise KeyError(key) diff --git a/source/tests/common/test_model_format_utils.py b/source/tests/common/test_model_format_utils.py index 57564df1ae..cb85fd2bb2 100644 --- a/source/tests/common/test_model_format_utils.py +++ b/source/tests/common/test_model_format_utils.py @@ -471,3 +471,22 @@ def test_self_exception( with self.assertRaises(ValueError) as context: ret0 = ifn0(dd[0], atype, fparam=ifp, aparam=iap) self.assertIn("input aparam", context.exception) + + def test_get_set(self): + ifn0 = InvarFitting( + "energy", + self.nt, + 3, + 1, + ) + rng = np.random.default_rng() + foo = rng.normal([3, 4]) + for ii in [ + "bias_atom_e", + "fparam_avg", + "fparam_inv_std", + "aparam_avg", + "aparam_inv_std", + ]: + ifn0[ii] = foo + np.testing.assert_allclose(foo, ifn0[ii]) diff --git a/source/tests/pt/test_ener_fitting.py b/source/tests/pt/test_ener_fitting.py index f0d6bdb932..eece8447df 100644 --- a/source/tests/pt/test_ener_fitting.py +++ b/source/tests/pt/test_ener_fitting.py @@ -160,3 +160,22 @@ def test_jit( use_tebd=(not distinguish_types), ).to(env.DEVICE) torch.jit.script(ft0) + + def test_get_set(self): + ifn0 = InvarFitting( + "energy", + self.nt, + 3, + 1, + ) + rng = np.random.default_rng() + foo = rng.normal([3, 4]) + for ii in [ + "bias_atom_e", + "fparam_avg", + "fparam_inv_std", + "aparam_avg", + "aparam_inv_std", + ]: + ifn0[ii] = torch.tensor(foo, dtype=dtype, device=env.DEVICE) + np.testing.assert_allclose(foo, ifn0[ii].detach().cpu().numpy()) From 73b588d12c8660d7c056276896b2bfd275374f09 Mon Sep 17 00:00:00 2001 From: Han Wang Date: Tue, 30 Jan 2024 10:04:18 +0800 Subject: [PATCH 13/14] reverse map for dtypes. add uts --- deepmd/pt/utils/utils.py | 36 +++++++++++++++++++++-------------- source/tests/pt/test_utils.py | 31 ++++++++++++++++++++++++++++++ 2 files changed, 53 insertions(+), 14 deletions(-) create mode 100644 source/tests/pt/test_utils.py diff --git a/deepmd/pt/utils/utils.py b/deepmd/pt/utils/utils.py index 516cbbdba6..e83e12f608 100644 --- a/deepmd/pt/utils/utils.py +++ b/deepmd/pt/utils/utils.py @@ -54,22 +54,30 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: def to_numpy_array( xx: torch.Tensor, ) -> np.ndarray: - if xx is not None: - prec = [key for key, value in PT_PRECISION_DICT.items() if value == xx.dtype] - if len(prec) == 0: - raise ValueError(f"unknown precision {xx.dtype}") - else: - prec = NP_PRECISION_DICT[prec[0]] - return xx.detach().cpu().numpy().astype(prec) if xx is not None else None + if xx is None: + return None + assert xx is not None + # Create a reverse mapping of PT_PRECISION_DICT + reverse_precision_dict = {v: k for k, v in PT_PRECISION_DICT.items()} + # Use the reverse mapping to find keys with the desired value + prec = reverse_precision_dict.get(xx.dtype, None) + prec = NP_PRECISION_DICT.get(prec, None) + if prec is None: + raise ValueError(f"unknown precision {xx.dtype}") + return xx.detach().cpu().numpy().astype(prec) def to_torch_tensor( xx: np.ndarray, ) -> torch.Tensor: - if xx is not None: - prec = [key for key, value in NP_PRECISION_DICT.items() if value == xx.dtype] - if len(prec) == 0: - raise ValueError(f"unknown precision {xx.dtype}") - else: - prec = PT_PRECISION_DICT[prec[0]] - return torch.tensor(xx, dtype=prec, device=DEVICE) if xx is not None else None + if xx is None: + return None + assert xx is not None + # Create a reverse mapping of NP_PRECISION_DICT + reverse_precision_dict = {v: k for k, v in NP_PRECISION_DICT.items()} + # Use the reverse mapping to find keys with the desired value + prec = reverse_precision_dict.get(type(xx.flat[0]), None) + prec = PT_PRECISION_DICT.get(prec, None) + if prec is None: + raise ValueError(f"unknown precision {xx.dtype}") + return torch.tensor(xx, dtype=prec, device=DEVICE) diff --git a/source/tests/pt/test_utils.py b/source/tests/pt/test_utils.py new file mode 100644 index 0000000000..9c9a9479ad --- /dev/null +++ b/source/tests/pt/test_utils.py @@ -0,0 +1,31 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import unittest + +import numpy as np +import torch + +from deepmd.pt.utils.utils import ( + to_numpy_array, + to_torch_tensor, +) + + +class TestCvt(unittest.TestCase): + def test_to_numpy(self): + rng = np.random.default_rng() + foo = rng.normal([3, 4]) + for ptp, npp in zip( + [torch.float16, torch.float32, torch.float64], + [np.float16, np.float32, np.float64], + ): + foo = foo.astype(npp) + bar = to_torch_tensor(foo) + self.assertEqual(bar.dtype, ptp) + onk = to_numpy_array(bar) + self.assertEqual(onk.dtype, npp) + with self.assertRaises(ValueError) as ee: + foo = foo.astype(np.int32) + bar = to_torch_tensor(foo) + with self.assertRaises(ValueError) as ee: + bar = to_torch_tensor(foo) + bar = to_numpy_array(bar.int()) From ce3116be3ba14977fd975e7f15c2795f8b7b219f Mon Sep 17 00:00:00 2001 From: Han Wang Date: Tue, 30 Jan 2024 10:38:46 +0800 Subject: [PATCH 14/14] rm unused vars --- deepmd/pt/model/task/ener.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/deepmd/pt/model/task/ener.py b/deepmd/pt/model/task/ener.py index 1d241469c2..91ccd03d9a 100644 --- a/deepmd/pt/model/task/ener.py +++ b/deepmd/pt/model/task/ener.py @@ -272,8 +272,6 @@ def forward( """ xx = descriptor nf, nloc, nd = xx.shape - dtype = descriptor.dtype - device = env.DEVICE # NOTICE in tests/pt/test_model.py # it happens that the user directly access the data memeber self.bias_atom_e # and set it to a wrong shape!