Merge branch 'deepmodeling:devel' into devel

deepmd/pt/infer/deep_eval.py

@@ -393,7 +393,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,

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]:

deepmd/pt/loss/ener_spin.py

@@ -63,25 +63,30 @@ 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
             The local atom number.
 
         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]:

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
 

deepmd/pt/loss/tensor.py

@@ -63,25 +63,30 @@ 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
             The local atom number.
 
         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]:

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(

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,

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"]

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,
         ),

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()
         )

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,