pd: support dpa3 with paddle backend

.pre-commit-config.yaml

@@ -65,13 +65,13 @@ repos:
       - id: clang-format
         exclude: ^(source/3rdparty|source/lib/src/gpu/cudart/.+\.inc|.+\.ipynb$|.+\.json$)
   # markdown, yaml, CSS, javascript
-  - repo: https://github.com/pre-commit/mirrors-prettier
-    rev: v4.0.0-alpha.8
-    hooks:
-      - id: prettier
-        types_or: [markdown, yaml, css]
-        # workflow files cannot be modified by pre-commit.ci
-        exclude: ^(source/3rdparty|\.github/workflows|\.clang-format)
+  # - repo: https://github.com/pre-commit/mirrors-prettier
+  #   rev: v4.0.0-alpha.8
+  #   hooks:
+  #     - id: prettier
+  #       types_or: [markdown, yaml, css]
+  #       # workflow files cannot be modified by pre-commit.ci
+  #       exclude: ^(source/3rdparty|\.github/workflows|\.clang-format)
   # Shell
   - repo: https://github.com/scop/pre-commit-shfmt
     rev: v3.11.0-1
@@ -83,25 +83,25 @@ repos:
     hooks:
       - id: cmake-format
       #- id: cmake-lint
-  - repo: https://github.com/njzjz/mirrors-bibtex-tidy
-    rev: v1.13.0
-    hooks:
-      - id: bibtex-tidy
-        args:
-          - --curly
-          - --numeric
-          - --align=13
-          - --blank-lines
-          # disable sort: the order of keys and fields has explict meanings
-          #- --sort=key
-          - --duplicates=key,doi,citation,abstract
-          - --merge=combine
-          #- --sort-fields
-          #- --strip-comments
-          - --trailing-commas
-          - --encode-urls
-          - --remove-empty-fields
-          - --wrap=80
+  # - repo: https://github.com/njzjz/mirrors-bibtex-tidy
+  #   rev: v1.13.0
+  #   hooks:
+  #     - id: bibtex-tidy
+  #       args:
+  #         - --curly
+  #         - --numeric
+  #         - --align=13
+  #         - --blank-lines
+  #         # disable sort: the order of keys and fields has explict meanings
+  #         #- --sort=key
+  #         - --duplicates=key,doi,citation,abstract
+  #         - --merge=combine
+  #         #- --sort-fields
+  #         #- --strip-comments
+  #         - --trailing-commas
+  #         - --encode-urls
+  #         - --remove-empty-fields
+  #         - --wrap=80
   # license header
   - repo: https://github.com/Lucas-C/pre-commit-hooks
     rev: v1.5.5

deepmd/infer/deep_eval.py

@@ -275,10 +275,6 @@ def get_has_spin(self) -> bool:
         """Check if the model has spin atom types."""
         return False
 
-    def get_has_hessian(self):
-        """Check if the model has hessian."""
-        return False
-
     def get_var_name(self) -> str:
         """Get the name of the fitting property."""
         raise NotImplementedError

deepmd/pd/loss/ener.py

@@ -56,7 +56,7 @@ def __init__(
         use_huber=False,
         huber_delta=0.01,
         **kwargs,
-    ):
+    ) -> None:
         r"""Construct a layer to compute loss on energy, force and virial.
 
         Parameters
@@ -225,28 +225,26 @@ def forward(self, input_dict, model, label, natoms, learning_rate, mae=False):
                 )
                 # more_loss['log_keys'].append('rmse_e')
             else:  # use l1 and for all atoms
+                energy_pred = energy_pred * atom_norm
+                energy_label = energy_label * atom_norm
                 l1_ener_loss = F.l1_loss(
                     energy_pred.reshape([-1]),
                     energy_label.reshape([-1]),
-                    reduction="sum",
+                    reduction="mean",
                 )
                 loss += pref_e * l1_ener_loss
                 more_loss["mae_e"] = self.display_if_exist(
-                    F.l1_loss(
-                        energy_pred.reshape([-1]),
-                        energy_label.reshape([-1]),
-                        reduction="mean",
-                    ).detach(),
+                    l1_ener_loss.detach(),
                     find_energy,
                 )
                 # more_loss['log_keys'].append('rmse_e')
-            if mae:
-                mae_e = paddle.mean(paddle.abs(energy_pred - energy_label)) * atom_norm
-                more_loss["mae_e"] = self.display_if_exist(mae_e.detach(), find_energy)
-                mae_e_all = paddle.mean(paddle.abs(energy_pred - energy_label))
-                more_loss["mae_e_all"] = self.display_if_exist(
-                    mae_e_all.detach(), find_energy
-                )
+            # if mae:
+            #     mae_e = torch.mean(torch.abs(energy_pred - energy_label)) * atom_norm
+            #     more_loss["mae_e"] = self.display_if_exist(mae_e.detach(), find_energy)
+            #     mae_e_all = torch.mean(torch.abs(energy_pred - energy_label))
+            #     more_loss["mae_e_all"] = self.display_if_exist(
+            #         mae_e_all.detach(), find_energy
+            #     )
 
         if (
             (self.has_f or self.has_pf or self.relative_f or self.has_gf)
@@ -287,17 +285,17 @@ def forward(self, input_dict, model, label, natoms, learning_rate, mae=False):
                         rmse_f.detach(), find_force
                     )
                 else:
-                    l1_force_loss = F.l1_loss(force_label, force_pred, reduction="none")
+                    l1_force_loss = F.l1_loss(force_label, force_pred, reduction="mean")
                     more_loss["mae_f"] = self.display_if_exist(
-                        l1_force_loss.mean().detach(), find_force
+                        l1_force_loss.detach(), find_force
                     )
-                    l1_force_loss = l1_force_loss.sum(-1).mean(-1).sum()
+                    # l1_force_loss = l1_force_loss.sum(-1).mean(-1).sum()
                     loss += (pref_f * l1_force_loss).to(GLOBAL_PD_FLOAT_PRECISION)
-                if mae:
-                    mae_f = paddle.mean(paddle.abs(diff_f))
-                    more_loss["mae_f"] = self.display_if_exist(
-                        mae_f.detach(), find_force
-                    )
+                # if mae:
+                #     mae_f = torch.mean(torch.abs(diff_f))
+                #     more_loss["mae_f"] = self.display_if_exist(
+                #         mae_f.detach(), find_force
+                #     )
 
             if self.has_pf and "atom_pref" in label:
                 atom_pref = label["atom_pref"]
@@ -324,20 +322,19 @@ def forward(self, input_dict, model, label, natoms, learning_rate, mae=False):
                 drdq_reshape = drdq.reshape(
                     [-1, natoms * 3, self.numb_generalized_coord]
                 )
+                gen_force_label = paddle.einsum(
+                    "bij,bi->bj", drdq_reshape, force_label_reshape_nframes
+                )
+                # gen_force_label = (
+                #     drdq_reshape * force_label_reshape_nframes.unsqueeze(-1)
+                # ).sum([-2])
 
-                # gen_force_label = paddle.einsum(
-                #     "bij,bi->bj", drdq_reshape, force_label_reshape_nframes
-                # )
-                gen_force_label = (
-                    drdq_reshape * force_label_reshape_nframes.unsqueeze(-1)
-                ).sum([-2])
-
-                # gen_force = paddle.einsum(
-                #     "bij,bi->bj", drdq_reshape, force_reshape_nframes
-                # )
-                gen_force = (drdq_reshape * force_reshape_nframes.unsqueeze(-1)).sum(
-                    [-2]
+                gen_force = paddle.einsum(
+                    "bij,bi->bj", drdq_reshape, force_reshape_nframes
                 )
+                # gen_force = (drdq_reshape * force_reshape_nframes.unsqueeze(-1)).sum(
+                #     [-2]
+                # )
 
                 diff_gen_force = gen_force_label - gen_force
                 l2_gen_force_loss = paddle.square(diff_gen_force).mean()
@@ -354,11 +351,33 @@ def forward(self, input_dict, model, label, natoms, learning_rate, mae=False):
         if self.has_v and "virial" in model_pred and "virial" in label:
             find_virial = label.get("find_virial", 0.0)
             pref_v = pref_v * find_virial
+            virial_label = label["virial"]
+            virial_pred = model_pred["virial"].reshape([-1, 9])
             diff_v = label["virial"] - model_pred["virial"].reshape([-1, 9])
-            l2_virial_loss = paddle.mean(paddle.square(diff_v))
-            if not self.inference:
-                more_loss["l2_virial_loss"] = self.display_if_exist(
-                    l2_virial_loss.detach(), find_virial
+            if not self.use_l1_all:
+                l2_virial_loss = paddle.mean(paddle.square(diff_v))
+                if not self.inference:
+                    more_loss["l2_virial_loss"] = self.display_if_exist(
+                        l2_virial_loss.detach(), find_virial
+                    )
+                if not self.huber:
+                    loss += atom_norm * (pref_v * l2_virial_loss)
+                else:
+                    if self.torch_huber:
+                        l_huber_loss = self.huber_loss(
+                            atom_norm * model_pred["virial"],
+                            atom_norm * label["virial"],
+                        )
+                    else:
+                        l_huber_loss = custom_huber_loss(
+                            atom_norm * model_pred["virial"].reshape([-1]),
+                            atom_norm * label["virial"].reshape([-1]),
+                            delta=self.huber_delta,
+                        )
+                    loss += pref_v * l_huber_loss
+                rmse_v = l2_virial_loss.sqrt() * atom_norm
+                more_loss["rmse_v"] = self.display_if_exist(
+                    rmse_v.detach(), find_virial
                 )
             if not self.use_huber:
                 loss += atom_norm * (pref_v * l2_virial_loss)
@@ -534,3 +553,75 @@ def deserialize(cls, data: dict) -> "TaskLoss":
         check_version_compatibility(data.pop("@version"), 2, 1)
         data.pop("@class")
         return cls(**data)
+
+
+class EnergyHessianStdLoss(EnergyStdLoss):
+    def __init__(
+        self,
+        start_pref_h=0.0,
+        limit_pref_h=0.0,
+        **kwargs,
+    ):
+        r"""Enable the layer to compute loss on hessian.
+
+        Parameters
+        ----------
+        start_pref_h : float
+            The prefactor of hessian loss at the start of the training.
+        limit_pref_h : float
+            The prefactor of hessian loss at the end of the training.
+        **kwargs
+            Other keyword arguments.
+        """
+        super().__init__(**kwargs)
+        self.has_h = (start_pref_h != 0.0 and limit_pref_h != 0.0) or self.inference
+
+        self.start_pref_h = start_pref_h
+        self.limit_pref_h = limit_pref_h
+
+    def forward(self, input_dict, model, label, natoms, learning_rate, mae=False):
+        model_pred, loss, more_loss = super().forward(
+            input_dict, model, label, natoms, learning_rate, mae=mae
+        )
+        coef = learning_rate / self.starter_learning_rate
+        pref_h = self.limit_pref_h + (self.start_pref_h - self.limit_pref_h) * coef
+
+        if self.has_h and "hessian" in model_pred and "hessian" in label:
+            find_hessian = label.get("find_hessian", 0.0)
+            pref_h = pref_h * find_hessian
+            diff_h = label["hessian"].reshape(
+                [-1],
+            ) - model_pred["hessian"].reshape(
+                [-1],
+            )
+            l2_hessian_loss = paddle.mean(paddle.square(diff_h))
+            if not self.inference:
+                more_loss["l2_hessian_loss"] = self.display_if_exist(
+                    l2_hessian_loss.detach(), find_hessian
+                )
+            loss += pref_h * l2_hessian_loss
+            rmse_h = l2_hessian_loss.sqrt()
+            more_loss["rmse_h"] = self.display_if_exist(rmse_h.detach(), find_hessian)
+            if mae:
+                mae_h = paddle.mean(paddle.abs(diff_h))
+                more_loss["mae_h"] = self.display_if_exist(mae_h.detach(), find_hessian)
+
+        if not self.inference:
+            more_loss["rmse"] = paddle.sqrt(loss.detach())
+        return model_pred, loss, more_loss
+
+    @property
+    def label_requirement(self) -> list[DataRequirementItem]:
+        """Add hessian label requirement needed for this loss calculation."""
+        label_requirement = super().label_requirement
+        if self.has_h:
+            label_requirement.append(
+                DataRequirementItem(
+                    "hessian",
+                    ndof=1,  # 9=3*3 --> 3N*3N=ndof*natoms*natoms
+                    atomic=True,
+                    must=False,
+                    high_prec=False,
+                )
+            )
+        return label_requirement

deepmd/pd/model/descriptor/__init__.py

@@ -12,6 +12,9 @@
 from .dpa2 import (
     DescrptDPA2,
 )
+from .dpa3 import (
+    DescrptDPA3,
+)
 from .env_mat import (
     prod_env_mat,
 )
@@ -39,6 +42,7 @@
     "DescrptBlockSeTTebd",
     "DescrptDPA1",
     "DescrptDPA2",
+    "DescrptDPA3",
     "DescrptSeA",
     "DescrptSeAttenV2",
     "DescrptSeTTebd",