feat: JAX training

deepmd/backend/jax.py

@@ -60,7 +60,11 @@
         Callable[[Namespace], None]
             The entry point hook of the backend.
         """
-        raise NotImplementedError
+        from deepmd.jax.entrypoints.main import (
+            main,
+        )
+
+        return main
 
     @property
     def deep_eval(self) -> type["DeepEvalBackend"]:

deepmd/dpmodel/descriptor/dpa1.py

@@ -781,6 +781,7 @@ def __init__(
         self.mean = np.zeros(wanted_shape, dtype=PRECISION_DICT[self.precision])
         self.stddev = np.ones(wanted_shape, dtype=PRECISION_DICT[self.precision])
         self.orig_sel = self.sel
+        self.ndescrpt = self.nnei * 4
 
     def get_rcut(self) -> float:
         """Returns the cut-off radius."""

deepmd/dpmodel/fitting/ener_fitting.py

@@ -6,6 +6,8 @@
     Union,
 )
 
+import numpy as np
+
 from deepmd.dpmodel.common import (
     DEFAULT_PRECISION,
 )
@@ -17,6 +19,10 @@
     from deepmd.dpmodel.fitting.general_fitting import (
         GeneralFitting,
     )
+
+from deepmd.utils.out_stat import (
+    compute_stats_from_redu,
+)
 from deepmd.utils.version import (
     check_version_compatibility,
 )
@@ -86,3 +92,69 @@
             **super().serialize(),
             "type": "ener",
         }
+
+    def compute_output_stats(self, all_stat: dict, mixed_type: bool = False) -> None:
+        """Compute the output statistics.
+
+        Parameters
+        ----------
+        all_stat
+            must have the following components:
+            all_stat['energy'] of shape n_sys x n_batch x n_frame
+            can be prepared by model.make_stat_input
+        mixed_type
+            Whether to perform the mixed_type mode.
+            If True, the input data has the mixed_type format (see doc/model/train_se_atten.md),
+            in which frames in a system may have different natoms_vec(s), with the same nloc.
+        """
+        self.bias_atom_e = self._compute_output_stats(
+            all_stat, rcond=self.rcond, mixed_type=mixed_type
+        )
+
+    def _compute_output_stats(self, all_stat, rcond=1e-3, mixed_type=False):
+        data = all_stat["energy"]
+        # data[sys_idx][batch_idx][frame_idx]
+        sys_ener = []
+        for ss in range(len(data)):
+            sys_data = []
+            for ii in range(len(data[ss])):
+                for jj in range(len(data[ss][ii])):
+                    sys_data.append(data[ss][ii][jj])
+            sys_data = np.concatenate(sys_data)
+            sys_ener.append(np.average(sys_data))
+        sys_ener = np.array(sys_ener)
+        sys_tynatom = []
+        if mixed_type:
+            data = all_stat["real_natoms_vec"]
+            nsys = len(data)
+            for ss in range(len(data)):
+                tmp_tynatom = []
+                for ii in range(len(data[ss])):
+                    for jj in range(len(data[ss][ii])):
+                        tmp_tynatom.append(data[ss][ii][jj].astype(np.float64))
+                tmp_tynatom = np.average(np.array(tmp_tynatom), axis=0)
+                sys_tynatom.append(tmp_tynatom)
+        else:
+            data = all_stat["natoms_vec"]
+            nsys = len(data)
+            for ss in range(len(data)):
+                sys_tynatom.append(data[ss][0].astype(np.float64))
+        sys_tynatom = np.array(sys_tynatom)
+        sys_tynatom = np.reshape(sys_tynatom, [nsys, -1])
+        sys_tynatom = sys_tynatom[:, 2:]
+        if len(self.atom_ener) > 0:
+            # Atomic energies stats are incorrect if atomic energies are assigned.
+            # In this situation, we directly use these assigned energies instead of computing stats.
+            # This will make the loss decrease quickly
+            assigned_atom_ener = np.array(
+                [ee if ee is not None else np.nan for ee in self.atom_ener_v]
+            )
+        else:
+            assigned_atom_ener = None
+        energy_shift, _ = compute_stats_from_redu(
+            sys_ener.reshape(-1, 1),
+            sys_tynatom,
+            assigned_bias=assigned_atom_ener,
+            rcond=rcond,
+        )
+        return energy_shift.ravel()

deepmd/dpmodel/loss/ener.py

@@ -177,7 +177,9 @@
                     delta=self.huber_delta,
                 )
                 loss += pref_e * l_huber_loss
-            more_loss["rmse_e"] = self.display_if_exist(l2_ener_loss, find_energy)
+            more_loss["rmse_e"] = self.display_if_exist(
+                xp.sqrt(l2_ener_loss) * atom_norm_ener, find_energy
+            )
         if self.has_f:
             l2_force_loss = xp.mean(xp.square(diff_f))
             if not self.use_huber:
@@ -189,7 +191,9 @@
                     delta=self.huber_delta,
                 )
                 loss += pref_f * l_huber_loss
-            more_loss["rmse_f"] = self.display_if_exist(l2_force_loss, find_force)
+            more_loss["rmse_f"] = self.display_if_exist(
+                xp.sqrt(l2_force_loss), find_force
+            )
         if self.has_v:
             virial_reshape = xp.reshape(virial, [-1])
             virial_hat_reshape = xp.reshape(virial_hat, [-1])
@@ -381,3 +385,79 @@
         check_version_compatibility(data.pop("@version"), 2, 1)
         data.pop("@class")
         return cls(**data)
+
+
+class EnergyHessianLoss(EnergyLoss):
+    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.
+        """
+        EnergyLoss.__init__(self, **kwargs)
+        self.has_h = start_pref_h != 0.0 and limit_pref_h != 0.0
+
+        self.start_pref_h = start_pref_h
+        self.limit_pref_h = limit_pref_h
+
+    def call(
+        self,
+        learning_rate: float,
+        natoms: int,
+        model_dict: dict[str, np.ndarray],
+        label_dict: dict[str, np.ndarray],
+    ) -> dict[str, np.ndarray]:
+        """Calculate loss from model results and labeled results."""
+        loss, more_loss = EnergyLoss.call(
+            self, learning_rate, natoms, model_dict, label_dict
+        )
+        xp = array_api_compat.array_namespace(model_dict["energy"])
+        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 "energy_derv_r_derv_r" in model_dict
+            and "hessian" in label_dict
+        ):
+            find_hessian = label_dict.get("find_hessian", 0.0)
+            pref_h = pref_h * find_hessian
+            diff_h = label_dict["hessian"].reshape(
+                -1,
+            ) - model_dict["energy_derv_r_derv_r"].reshape(
+                -1,
+            )
+            l2_hessian_loss = xp.mean(xp.square(diff_h))
+            loss += pref_h * l2_hessian_loss
+            rmse_h = xp.sqrt(l2_hessian_loss)
+            more_loss["rmse_h"] = self.display_if_exist(rmse_h, find_hessian)
+
+        more_loss["rmse"] = xp.sqrt(loss)
+        return 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/dpmodel/utils/env_mat_stat.py

@@ -119,12 +119,15 @@ def iter(
                 "last_dim should be 1 for raial-only or 4 for full descriptor."
             )
         for system in data:
-            coord, atype, box, natoms = (
+            coord, atype, box = (
                 system["coord"],
                 system["atype"],
                 system["box"],
-                system["natoms"],
             )
+            coord = xp.reshape(coord, (coord.shape[0], -1, 3))  # (nframes, nloc, 3)
+            atype = xp.reshape(atype, (coord.shape[0], -1))  # (nframes, nloc)
+            if box is not None:
+                box = xp.reshape(box, (coord.shape[0], 3, 3))
             (
                 extended_coord,
                 extended_atype,

deepmd/dpmodel/utils/learning_rate.py

@@ -45,9 +45,8 @@ def __init__(
             self.decay_rate = decay_rate
         self.min_lr = stop_lr
 
-    def value(self, step) -> np.float64:
+    def value(self, step, xp=np) -> np.float64:
         """Get the learning rate at the given step."""
-        step_lr = self.start_lr * np.power(self.decay_rate, step // self.decay_steps)
-        if step_lr < self.min_lr:
-            step_lr = self.min_lr
+        step_lr = self.start_lr * xp.power(self.decay_rate, step // self.decay_steps)
+        step_lr = xp.clip(step_lr, self.min_lr, None)
         return step_lr

deepmd/jax/entrypoints/__init__.py

@@ -0,0 +1 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later

deepmd/jax/entrypoints/freeze.py

@@ -0,0 +1,36 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+from pathlib import (
+    Path,
+)
+
+from deepmd.jax.utils.serialization import (
+    deserialize_to_file,
+    serialize_from_file,
+)
+
+
+def freeze(
+    *,
+    checkpoint_folder: str,
+    output: str,
+    **kwargs,
+) -> None:
+    """Freeze the graph in supplied folder.
+
+    Parameters
+    ----------
+    checkpoint_folder : str
+        location of either the folder with checkpoint or the checkpoint prefix
+    output : str
+        output file name
+    **kwargs
+        other arguments
+    """
+    if (Path(checkpoint_folder) / "checkpoint").is_file():
+        checkpoint_meta = Path(checkpoint_folder) / "checkpoint"
+        checkpoint_folder = checkpoint_meta.read_text().strip()
+    if Path(checkpoint_folder).is_dir():
+        data = serialize_from_file(checkpoint_folder)
+        deserialize_to_file(output, data)
+    else:
+        raise FileNotFoundError(f"Checkpoint {checkpoint_folder} does not exist.")

deepmd/jax/entrypoints/main.py

@@ -0,0 +1,67 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""DeePMD-Kit entry point module."""
+
+import argparse
+from pathlib import (
+    Path,
+)
+from typing import (
+    Optional,
+    Union,
+)
+
+from deepmd.backend.suffix import (
+    format_model_suffix,
+)
+from deepmd.jax.entrypoints.freeze import (
+    freeze,
+)
+from deepmd.jax.entrypoints.train import (
+    train,
+)
+from deepmd.loggers.loggers import (
+    set_log_handles,
+)
+from deepmd.main import (
+    parse_args,
+)
+
+__all__ = ["main"]
+
+
+def main(args: Optional[Union[list[str], argparse.Namespace]] = None) -> None:
+    """DeePMD-Kit entry point.
+
+    Parameters
+    ----------
+    args : list[str] or argparse.Namespace, optional
+        list of command line arguments, used to avoid calling from the subprocess,
+        as it is quite slow to import tensorflow; if Namespace is given, it will
+        be used directly
+
+    Raises
+    ------
+    RuntimeError
+        if no command was input
+    """
+    if not isinstance(args, argparse.Namespace):
+        args = parse_args(args=args)
+
+    dict_args = vars(args)
+    set_log_handles(
+        args.log_level,
+        Path(args.log_path) if args.log_path else None,
+        mpi_log=None,
+    )
+
+    if args.command == "train":
+        train(**dict_args)
+    elif args.command == "freeze":
+        dict_args["output"] = format_model_suffix(
+            dict_args["output"], preferred_backend=args.backend, strict_prefer=True
+        )
+        freeze(**dict_args)
+    elif args.command is None:
+        pass
+    else:
+        raise RuntimeError(f"unknown command {args.command}")