feat(pt): add cosine annealing lr scheduler

deepmd/dpmodel/utils/learning_rate.py

@@ -55,3 +55,34 @@ def value(self, step: int) -> np.float64:
         if step_lr < self.min_lr:
             step_lr = self.min_lr
         return step_lr
+
+
+class LearningRateCosine:
+    def __init__(
+        self,
+        start_lr: float,
+        stop_lr: float,
+        stop_steps: int,
+        **kwargs: Any,
+    ) -> None:
+        """
+        Construct a cosine-annealed learning rate.
+
+        Parameters
+        ----------
+        start_lr : float
+            The learning rate at the start of the training.
+        stop_lr : float
+            The desired learning rate at the end of the training.
+        stop_steps : int
+            The total training steps for learning rate scheduler.
+        """
+        self.start_lr = start_lr
+        self.stop_lr = stop_lr
+        self.stop_steps = max(1, stop_steps)
+
+    def value(self, step: int) -> np.float64:
+        """Get the learning rate at the given step."""
+        clamped_step = min(step, self.stop_steps)
+        cosine = 0.5 * (1.0 + np.cos(np.pi * clamped_step / self.stop_steps))
+        return np.float64(self.stop_lr + (self.start_lr - self.stop_lr) * cosine)

deepmd/pt/train/training.py

@@ -63,6 +63,7 @@
     SAMPLER_RECORD,
 )
 from deepmd.pt.utils.learning_rate import (
+    LearningRateCosine,
     LearningRateExp,
 )
 from deepmd.pt.utils.stat import (
@@ -273,13 +274,17 @@ def get_sample() -> Any:
                     _stat_file_path.root.close()
             return get_sample
 
-        def get_lr(lr_params: dict[str, Any]) -> LearningRateExp:
-            assert lr_params.get("type", "exp") == "exp", (
-                "Only learning rate `exp` is supported!"
-            )
-            lr_params["stop_steps"] = self.num_steps - self.warmup_steps
-            lr_exp = LearningRateExp(**lr_params)
-            return lr_exp
+        def get_lr(
+            lr_params: dict[str, Any],
+        ) -> LearningRateExp | LearningRateCosine:
+            lr_type = lr_params.get("type", "exp")
+            lr_config = lr_params.copy()
+            lr_config["stop_steps"] = self.num_steps - self.warmup_steps
+            if lr_type == "exp":
+                return LearningRateExp(**lr_config)
+            if lr_type == "cosine":
+                return LearningRateCosine(**lr_config)
+            raise ValueError(f"Unsupported learning rate type '{lr_type}'")
 
         # Optimizer
         if self.multi_task and training_params.get("optim_dict", None) is not None:

deepmd/pt/utils/learning_rate.py

@@ -1,8 +1,10 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
 from deepmd.dpmodel.utils.learning_rate import (
+    LearningRateCosine,
     LearningRateExp,
 )
 
 __all__ = [
+    "LearningRateCosine",
     "LearningRateExp",
 ]

deepmd/utils/argcheck.py

@@ -2509,12 +2509,31 @@ def learning_rate_exp() -> list[Argument]:
     return args
 
 
+def learning_rate_cosine() -> list[Argument]:
+    doc_start_lr = "The learning rate at the start of the training."
+    doc_stop_lr = "The desired learning rate at the end of the training."
+
+    args = [
+        Argument("start_lr", float, optional=True, default=1e-3, doc=doc_start_lr),
+        Argument("stop_lr", float, optional=True, default=1e-7, doc=doc_stop_lr),
+    ]
+    return args
+
+
 def learning_rate_variant_type_args() -> Variant:
     doc_lr = "The type of the learning rate."
 
     return Variant(
         "type",
-        [Argument("exp", dict, learning_rate_exp())],
+        [
+            Argument("exp", dict, learning_rate_exp()),
+            Argument(
+                "cosine",
+                dict,
+                learning_rate_cosine(),
+                doc=doc_only_pt_supported,
+            ),
+        ],
         optional=True,
         default_tag="exp",
         doc=doc_lr,