Merge branch 'devel' into save-jax

Signed-off-by: Jinzhe Zeng <[email protected]>

.pre-commit-config.yaml

@@ -29,7 +29,7 @@ repos:
         exclude: ^source/3rdparty
   - repo: https://github.com/astral-sh/ruff-pre-commit
     # Ruff version.
-    rev: v0.6.9
+    rev: v0.7.0
     hooks:
       - id: ruff
         args: ["--fix"]
@@ -52,7 +52,7 @@ repos:
       - id: blacken-docs
   # C++
   - repo: https://github.com/pre-commit/mirrors-clang-format
-    rev: v19.1.1
+    rev: v19.1.2
     hooks:
       - id: clang-format
         exclude: ^(source/3rdparty|source/lib/src/gpu/cudart/.+\.inc|.+\.ipynb$)
@@ -66,7 +66,7 @@ repos:
         exclude: ^(source/3rdparty|\.github/workflows|\.clang-format)
   # Shell
   - repo: https://github.com/scop/pre-commit-shfmt
-    rev: v3.9.0-1
+    rev: v3.10.0-1
     hooks:
       - id: shfmt
   # CMake

deepmd/dpmodel/atomic_model/base_atomic_model.py

@@ -243,7 +243,8 @@ def serialize(self) -> dict:
 
     @classmethod
     def deserialize(cls, data: dict) -> "BaseAtomicModel":
-        data = copy.deepcopy(data)
+        # do not deep copy Descriptor and Fitting class
+        data = data.copy()
         variables = data.pop("@variables")
         obj = cls(**data)
         for kk in variables.keys():

deepmd/dpmodel/fitting/dipole_fitting.py

@@ -105,7 +105,6 @@ def __init__(
         r_differentiable: bool = True,
         c_differentiable: bool = True,
         type_map: Optional[list[str]] = None,
-        old_impl=False,
         seed: Optional[Union[int, list[int]]] = None,
     ):
         if tot_ener_zero:
@@ -141,7 +140,6 @@ def __init__(
             type_map=type_map,
             seed=seed,
         )
-        self.old_impl = False
 
     def _net_out_dim(self):
         """Set the FittingNet output dim."""
@@ -151,7 +149,6 @@ def serialize(self) -> dict:
         data = super().serialize()
         data["type"] = "dipole"
         data["embedding_width"] = self.embedding_width
-        data["old_impl"] = self.old_impl
         data["r_differentiable"] = self.r_differentiable
         data["c_differentiable"] = self.c_differentiable
         return data

deepmd/dpmodel/fitting/polarizability_fitting.py

@@ -107,7 +107,6 @@ def __init__(
         spin: Any = None,
         mixed_types: bool = False,
         exclude_types: list[int] = [],
-        old_impl: bool = False,
         fit_diag: bool = True,
         scale: Optional[list[float]] = None,
         shift_diag: bool = True,
@@ -165,7 +164,6 @@ def __init__(
             type_map=type_map,
             seed=seed,
         )
-        self.old_impl = False
 
     def _net_out_dim(self):
         """Set the FittingNet output dim."""
@@ -192,7 +190,6 @@ def serialize(self) -> dict:
         data["type"] = "polar"
         data["@version"] = 3
         data["embedding_width"] = self.embedding_width
-        data["old_impl"] = self.old_impl
         data["fit_diag"] = self.fit_diag
         data["shift_diag"] = self.shift_diag
         data["@variables"]["scale"] = self.scale

deepmd/dpmodel/utils/learning_rate.py

@@ -0,0 +1,53 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+import numpy as np
+
+
+class LearningRateExp:
+    def __init__(
+        self,
+        start_lr,
+        stop_lr,
+        decay_steps,
+        stop_steps,
+        decay_rate=None,
+        **kwargs,
+    ):
+        """
+        Construct an exponential-decayed learning rate.
+
+        Parameters
+        ----------
+        start_lr
+            The learning rate at the start of the training.
+        stop_lr
+            The desired learning rate at the end of the training.
+            When decay_rate is explicitly set, this value will serve as
+            the minimum learning rate during training. In other words,
+            if the learning rate decays below stop_lr, stop_lr will be applied instead.
+        decay_steps
+            The learning rate is decaying every this number of training steps.
+        stop_steps
+            The total training steps for learning rate scheduler.
+        decay_rate
+            The decay rate for the learning rate.
+            If provided, the decay rate will be set instead of
+            calculating it through interpolation between start_lr and stop_lr.
+        """
+        self.start_lr = start_lr
+        default_ds = 100 if stop_steps // 10 > 100 else stop_steps // 100 + 1
+        self.decay_steps = decay_steps
+        if self.decay_steps >= stop_steps:
+            self.decay_steps = default_ds
+        self.decay_rate = np.exp(
+            np.log(stop_lr / self.start_lr) / (stop_steps / self.decay_steps)
+        )
+        if decay_rate is not None:
+            self.decay_rate = decay_rate
+        self.min_lr = stop_lr
+
+    def value(self, step) -> 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
+        return step_lr

deepmd/pt/entrypoints/main.py

@@ -283,7 +283,9 @@ def train(
     # update init_model or init_frz_model config if necessary
     if (init_model is not None or init_frz_model is not None) and use_pretrain_script:
         if init_model is not None:
-            init_state_dict = torch.load(init_model, map_location=DEVICE)
+            init_state_dict = torch.load(
+                init_model, map_location=DEVICE, weights_only=True
+            )
             if "model" in init_state_dict:
                 init_state_dict = init_state_dict["model"]
             config["model"] = init_state_dict["_extra_state"]["model_params"]
@@ -380,7 +382,9 @@ def change_bias(
     output: Optional[str] = None,
 ):
     if input_file.endswith(".pt"):
-        old_state_dict = torch.load(input_file, map_location=env.DEVICE)
+        old_state_dict = torch.load(
+            input_file, map_location=env.DEVICE, weights_only=True
+        )
         model_state_dict = copy.deepcopy(old_state_dict.get("model", old_state_dict))
         model_params = model_state_dict["_extra_state"]["model_params"]
     elif input_file.endswith(".pth"):

deepmd/pt/infer/deep_eval.py

@@ -103,7 +103,9 @@ def __init__(
         self.output_def = output_def
         self.model_path = model_file
         if str(self.model_path).endswith(".pt"):
-            state_dict = torch.load(model_file, map_location=env.DEVICE)
+            state_dict = torch.load(
+                model_file, map_location=env.DEVICE, weights_only=True
+            )
             if "model" in state_dict:
                 state_dict = state_dict["model"]
             self.input_param = state_dict["_extra_state"]["model_params"]

deepmd/pt/infer/inference.py

@@ -34,7 +34,7 @@ def __init__(
         - config: The Dict-like configuration with training options.
         """
         # Model
-        state_dict = torch.load(model_ckpt, map_location=DEVICE)
+        state_dict = torch.load(model_ckpt, map_location=DEVICE, weights_only=True)
         if "model" in state_dict:
             state_dict = state_dict["model"]
         model_params = state_dict["_extra_state"]["model_params"]

deepmd/pt/model/descriptor/__init__.py

@@ -16,9 +16,6 @@
 from .env_mat import (
     prod_env_mat,
 )
-from .gaussian_lcc import (
-    DescrptGaussianLcc,
-)
 from .hybrid import (
     DescrptHybrid,
 )
@@ -59,6 +56,5 @@
     "DescrptDPA2",
     "DescrptHybrid",
     "prod_env_mat",
-    "DescrptGaussianLcc",
     "DescrptBlockRepformers",
 ]

deepmd/pt/model/descriptor/dpa1.py

@@ -245,7 +245,6 @@ def __init__(
         # not implemented
         spin=None,
         type: Optional[str] = None,
-        old_impl: bool = False,
     ):
         super().__init__()
         # Ensure compatibility with the deprecated stripped_type_embedding option.
@@ -290,7 +289,6 @@ def __init__(
             trainable_ln=trainable_ln,
             ln_eps=ln_eps,
             seed=child_seed(seed, 1),
-            old_impl=old_impl,
         )
         self.use_econf_tebd = use_econf_tebd
         self.use_tebd_bias = use_tebd_bias

deepmd/pt/model/descriptor/dpa2.py

@@ -92,7 +92,6 @@ def __init__(
         use_econf_tebd: bool = False,
         use_tebd_bias: bool = False,
         type_map: Optional[list[str]] = None,
-        old_impl: bool = False,
     ):
         r"""The DPA-2 descriptor. see https://arxiv.org/abs/2312.15492.
 
@@ -235,7 +234,6 @@ def init_subclass_params(sub_data, sub_class):
             g1_out_conv=self.repformer_args.g1_out_conv,
             g1_out_mlp=self.repformer_args.g1_out_mlp,
             seed=child_seed(seed, 1),
-            old_impl=old_impl,
         )
         self.rcsl_list = [
             (self.repformers.get_rcut(), self.repformers.get_nsel()),