From 9c861c24f1be3874b08fe36274d8f4afea5b89f1 Mon Sep 17 00:00:00 2001 From: Duo <50307526+iProzd@users.noreply.github.com> Date: Wed, 20 Mar 2024 13:12:58 +0800 Subject: [PATCH 1/3] pt: add index input for `use_spin` (#3456) It's convenient for spin multitask input file to handle different spin inputs. --- deepmd/pt/model/model/__init__.py | 8 ++++++++ deepmd/utils/argcheck.py | 11 ++++++++--- 2 files changed, 16 insertions(+), 3 deletions(-) diff --git a/deepmd/pt/model/model/__init__.py b/deepmd/pt/model/model/__init__.py index 7a2070e476..1675215d7b 100644 --- a/deepmd/pt/model/model/__init__.py +++ b/deepmd/pt/model/model/__init__.py @@ -14,6 +14,8 @@ import copy import json +import numpy as np + from deepmd.pt.model.atomic_model import ( DPAtomicModel, PairTabAtomicModel, @@ -57,6 +59,12 @@ def get_spin_model(model_params): model_params = copy.deepcopy(model_params) + if not model_params["spin"]["use_spin"] or isinstance( + model_params["spin"]["use_spin"][0], int + ): + use_spin = np.full(len(model_params["type_map"]), False) + use_spin[model_params["spin"]["use_spin"]] = True + model_params["spin"]["use_spin"] = use_spin.tolist() # include virtual spin and placeholder types model_params["type_map"] += [item + "_spin" for item in model_params["type_map"]] spin = Spin( diff --git a/deepmd/utils/argcheck.py b/deepmd/utils/argcheck.py index 7a8e0e98cb..564039ccd0 100644 --- a/deepmd/utils/argcheck.py +++ b/deepmd/utils/argcheck.py @@ -93,7 +93,12 @@ def type_embedding_args(): def spin_args(): - doc_use_spin = "Whether to use atomic spin model for each atom type" + doc_use_spin = ( + "Whether to use atomic spin model for each atom type. " + "List of boolean values with the shape of [ntypes] to specify which types use spin, " + f"or a list of integer values {doc_only_pt_supported} " + "to indicate the index of the type that uses spin." + ) doc_spin_norm = "The magnitude of atomic spin for each atom type with spin" doc_virtual_len = "The distance between virtual atom representing spin and its corresponding real atom for each atom type with spin" doc_virtual_scale = ( @@ -106,7 +111,7 @@ def spin_args(): ) return [ - Argument("use_spin", List[bool], doc=doc_use_spin), + Argument("use_spin", [List[bool], List[int]], doc=doc_use_spin), Argument( "spin_norm", List[float], @@ -121,7 +126,7 @@ def spin_args(): ), Argument( "virtual_scale", - List[float], + [List[float], float], optional=True, doc=doc_only_pt_supported + doc_virtual_scale, ), From 71ec631167875277515014d07c3b8ab27b5c4e38 Mon Sep 17 00:00:00 2001 From: Duo <50307526+iProzd@users.noreply.github.com> Date: Wed, 20 Mar 2024 13:47:50 +0800 Subject: [PATCH 2/3] pt: refactor loss (#3569) This PR updates the loss interface to allow for a more flexible design. It enables processing input tensors before feeding them into the model, such as denoising operations (fyi @Chengqian-Zhang ). Previously, this was done in the data loader, which was less intuitive and more confusing. Now, users can easily handle these tasks within the loss function itself, as demonstrated in similar implementations in uni-mol: https://github.com/dptech-corp/Uni-Mol/blob/main/unimol/unimol/losses/unimol.py#L20. --- deepmd/pt/loss/ener.py | 30 ++++++++++++++------- deepmd/pt/loss/ener_spin.py | 13 ++++++--- deepmd/pt/loss/loss.py | 2 +- deepmd/pt/loss/tensor.py | 13 ++++++--- deepmd/pt/train/training.py | 9 +++++-- deepmd/pt/train/wrapper.py | 19 ++++++++----- source/tests/pt/model/test_model.py | 41 ++++++++++++++--------------- source/tests/pt/test_loss.py | 18 ++++++++++--- 8 files changed, 92 insertions(+), 53 deletions(-) diff --git a/deepmd/pt/loss/ener.py b/deepmd/pt/loss/ener.py index f29c3231f1..edae53a771 100644 --- a/deepmd/pt/loss/ener.py +++ b/deepmd/pt/loss/ener.py @@ -90,20 +90,30 @@ def __init__( self.use_l1_all = use_l1_all self.inference = inference - def forward(self, model_pred, label, natoms, learning_rate, mae=False): - """Return loss on loss and force. + def forward(self, input_dict, model, label, natoms, learning_rate, mae=False): + """Return loss on energy and force. - Args: - - natoms: Tell atom count. - - p_energy: Predicted energy of all atoms. - - p_force: Predicted force per atom. - - l_energy: Actual energy of all atoms. - - l_force: Actual force per atom. + Parameters + ---------- + input_dict : dict[str, torch.Tensor] + Model inputs. + model : torch.nn.Module + Model to be used to output the predictions. + label : dict[str, torch.Tensor] + Labels. + natoms : int + The local atom number. Returns ------- - - loss: Loss to minimize. + model_pred: dict[str, torch.Tensor] + Model predictions. + loss: torch.Tensor + Loss for model to minimize. + more_loss: dict[str, torch.Tensor] + Other losses for display. """ + model_pred = model(**input_dict) coef = learning_rate / self.starter_learning_rate pref_e = self.limit_pref_e + (self.start_pref_e - self.limit_pref_e) * coef pref_f = self.limit_pref_f + (self.start_pref_f - self.limit_pref_f) * coef @@ -200,7 +210,7 @@ def forward(self, model_pred, label, natoms, learning_rate, mae=False): more_loss["mae_v"] = mae_v.detach() if not self.inference: more_loss["rmse"] = torch.sqrt(loss.detach()) - return loss, more_loss + return model_pred, loss, more_loss @property def label_requirement(self) -> List[DataRequirementItem]: diff --git a/deepmd/pt/loss/ener_spin.py b/deepmd/pt/loss/ener_spin.py index 1f55dcc5df..1f10e3cf5f 100644 --- a/deepmd/pt/loss/ener_spin.py +++ b/deepmd/pt/loss/ener_spin.py @@ -63,13 +63,15 @@ def __init__( self.use_l1_all = use_l1_all self.inference = inference - def forward(self, model_pred, label, natoms, learning_rate, mae=False): + def forward(self, input_dict, model, label, natoms, learning_rate, mae=False): """Return energy loss with magnetic labels. Parameters ---------- - model_pred : dict[str, torch.Tensor] - Model predictions. + input_dict : dict[str, torch.Tensor] + Model inputs. + model : torch.nn.Module + Model to be used to output the predictions. label : dict[str, torch.Tensor] Labels. natoms : int @@ -77,11 +79,14 @@ def forward(self, model_pred, label, natoms, learning_rate, mae=False): Returns ------- + model_pred: dict[str, torch.Tensor] + Model predictions. loss: torch.Tensor Loss for model to minimize. more_loss: dict[str, torch.Tensor] Other losses for display. """ + model_pred = model(**input_dict) coef = learning_rate / self.starter_learning_rate pref_e = self.limit_pref_e + (self.start_pref_e - self.limit_pref_e) * coef pref_fr = self.limit_pref_fr + (self.start_pref_fr - self.limit_pref_fr) * coef @@ -175,7 +180,7 @@ def forward(self, model_pred, label, natoms, learning_rate, mae=False): if not self.inference: more_loss["rmse"] = torch.sqrt(loss.detach()) - return loss, more_loss + return model_pred, loss, more_loss @property def label_requirement(self) -> List[DataRequirementItem]: diff --git a/deepmd/pt/loss/loss.py b/deepmd/pt/loss/loss.py index 925ff8f4ef..cc253424ca 100644 --- a/deepmd/pt/loss/loss.py +++ b/deepmd/pt/loss/loss.py @@ -19,7 +19,7 @@ def __init__(self, **kwargs): """Construct loss.""" super().__init__() - def forward(self, model_pred, label, natoms, learning_rate): + def forward(self, input_dict, model, label, natoms, learning_rate): """Return loss .""" raise NotImplementedError diff --git a/deepmd/pt/loss/tensor.py b/deepmd/pt/loss/tensor.py index ec7ef0b323..238e6a7796 100644 --- a/deepmd/pt/loss/tensor.py +++ b/deepmd/pt/loss/tensor.py @@ -63,13 +63,15 @@ def __init__( "Can not assian zero weight both to `pref` and `pref_atomic`" ) - def forward(self, model_pred, label, natoms, learning_rate=0.0, mae=False): + def forward(self, input_dict, model, label, natoms, learning_rate=0.0, mae=False): """Return loss on local and global tensors. Parameters ---------- - model_pred : dict[str, torch.Tensor] - Model predictions. + input_dict : dict[str, torch.Tensor] + Model inputs. + model : torch.nn.Module + Model to be used to output the predictions. label : dict[str, torch.Tensor] Labels. natoms : int @@ -77,11 +79,14 @@ def forward(self, model_pred, label, natoms, learning_rate=0.0, mae=False): Returns ------- + model_pred: dict[str, torch.Tensor] + Model predictions. loss: torch.Tensor Loss for model to minimize. more_loss: dict[str, torch.Tensor] Other losses for display. """ + model_pred = model(**input_dict) del learning_rate, mae loss = torch.zeros(1, dtype=env.GLOBAL_PT_FLOAT_PRECISION, device=env.DEVICE)[0] more_loss = {} @@ -133,7 +138,7 @@ def forward(self, model_pred, label, natoms, learning_rate=0.0, mae=False): loss += self.global_weight * l2_global_loss rmse_global = l2_global_loss.sqrt() / atom_num more_loss[f"rmse_global_{self.tensor_name}"] = rmse_global.detach() - return loss, more_loss + return model_pred, loss, more_loss @property def label_requirement(self) -> List[DataRequirementItem]: diff --git a/deepmd/pt/train/training.py b/deepmd/pt/train/training.py index 2056b9b305..9fd675a8f2 100644 --- a/deepmd/pt/train/training.py +++ b/deepmd/pt/train/training.py @@ -696,8 +696,13 @@ def step(_step_id, task_key="Default"): module = ( self.wrapper.module if dist.is_initialized() else self.wrapper ) - loss, more_loss = module.loss[task_key]( - model_pred, + + def fake_model(): + return model_pred + + _, loss, more_loss = module.loss[task_key]( + {}, + fake_model, label_dict, int(input_dict["atype"].shape[-1]), learning_rate=pref_lr, diff --git a/deepmd/pt/train/wrapper.py b/deepmd/pt/train/wrapper.py index 061cd777db..6bc7cdc87a 100644 --- a/deepmd/pt/train/wrapper.py +++ b/deepmd/pt/train/wrapper.py @@ -168,15 +168,20 @@ def forward( has_spin = has_spin() if has_spin: input_dict["spin"] = spin - model_pred = self.model[task_key](**input_dict) - natoms = atype.shape[-1] - if not self.inference_only and not inference_only: - loss, more_loss = self.loss[task_key]( - model_pred, label, natoms=natoms, learning_rate=cur_lr + + if self.inference_only or inference_only: + model_pred = self.model[task_key](**input_dict) + return model_pred, None, None + else: + natoms = atype.shape[-1] + model_pred, loss, more_loss = self.loss[task_key]( + input_dict, + self.model[task_key], + label, + natoms=natoms, + learning_rate=cur_lr, ) return model_pred, loss, more_loss - else: - return model_pred, None, None def set_extra_state(self, state: Dict): self.model_params = state["model_params"] diff --git a/source/tests/pt/model/test_model.py b/source/tests/pt/model/test_model.py index 5a30de7ac8..aa1c0dd969 100644 --- a/source/tests/pt/model/test_model.py +++ b/source/tests/pt/model/test_model.py @@ -338,34 +338,33 @@ def test_consistency(self): batch["natoms"] = torch.tensor( batch["natoms_vec"], device=batch["coord"].device ).unsqueeze(0) - model_predict = my_model( - batch["coord"].to(env.DEVICE), - batch["atype"].to(env.DEVICE), - batch["box"].to(env.DEVICE), - do_atomic_virial=True, - ) - model_predict_1 = my_model( - batch["coord"].to(env.DEVICE), - batch["atype"].to(env.DEVICE), - batch["box"].to(env.DEVICE), - do_atomic_virial=False, + model_input = { + "coord": batch["coord"].to(env.DEVICE), + "atype": batch["atype"].to(env.DEVICE), + "box": batch["box"].to(env.DEVICE), + "do_atomic_virial": True, + } + model_input_1 = { + "coord": batch["coord"].to(env.DEVICE), + "atype": batch["atype"].to(env.DEVICE), + "box": batch["box"].to(env.DEVICE), + "do_atomic_virial": False, + } + label = { + "energy": batch["energy"].to(env.DEVICE), + "force": batch["force"].to(env.DEVICE), + } + cur_lr = my_lr.value(self.wanted_step) + model_predict, loss, _ = my_loss( + model_input, my_model, label, int(batch["natoms"][0, 0]), cur_lr ) + model_predict_1 = my_model(**model_input_1) p_energy, p_force, p_virial, p_atomic_virial = ( model_predict["energy"], model_predict["force"], model_predict["virial"], model_predict["atom_virial"], ) - cur_lr = my_lr.value(self.wanted_step) - model_pred = { - "energy": p_energy, - "force": p_force, - } - label = { - "energy": batch["energy"].to(env.DEVICE), - "force": batch["force"].to(env.DEVICE), - } - loss, _ = my_loss(model_pred, label, int(batch["natoms"][0, 0]), cur_lr) np.testing.assert_allclose( head_dict["energy"], p_energy.view(-1).cpu().detach().numpy() ) diff --git a/source/tests/pt/test_loss.py b/source/tests/pt/test_loss.py index dddc9af219..2abb22c2a9 100644 --- a/source/tests/pt/test_loss.py +++ b/source/tests/pt/test_loss.py @@ -171,8 +171,13 @@ def test_consistency(self): self.start_pref_v, self.limit_pref_v, ) - my_loss, my_more_loss = mine( - self.model_pred, + + def fake_model(): + return self.model_pred + + _, my_loss, my_more_loss = mine( + {}, + fake_model, self.label, self.nloc, self.cur_lr, @@ -345,8 +350,13 @@ def test_consistency(self): self.start_pref_fm, self.limit_pref_fm, ) - my_loss, my_more_loss = mine( - self.model_pred, + + def fake_model(): + return self.model_pred + + _, my_loss, my_more_loss = mine( + {}, + fake_model, self.label, self.nloc_tf, # use tf natoms pref self.cur_lr, From 5aa1b89901a31d8fd6fe55b33dcd9d3fee8f3d8d Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Wed, 20 Mar 2024 01:48:15 -0400 Subject: [PATCH 3/3] fix(pt): fix a typo in DeepEval to check do_atomic_virial (#3570) Signed-off-by: Jinzhe Zeng --- deepmd/pt/infer/deep_eval.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/deepmd/pt/infer/deep_eval.py b/deepmd/pt/infer/deep_eval.py index b8031993c0..f46d5fce49 100644 --- a/deepmd/pt/infer/deep_eval.py +++ b/deepmd/pt/infer/deep_eval.py @@ -388,7 +388,7 @@ def _eval_model( else: aparam_input = None do_atomic_virial = any( - x.category == OutputVariableCategory.DERV_C_REDU for x in request_defs + x.category == OutputVariableCategory.DERV_C for x in request_defs ) batch_output = model( coord_input,