add linear lr

Author: iProzd

deepmd/dpmodel/utils/learning_rate.py

@@ -109,3 +109,95 @@ def value(self, step) -> np.float64:
             return self.start_lr - decay_rate * (
                 step - self.decay_start_rate * self.stop_steps
             )
+
+
+class LearningRateLinear:
+    def __init__(
+        self,
+        start_lr: float,
+        stop_steps: int,
+        decay_steps: int,
+        start_factor: float = 1.0,
+        end_factor: float = 1.0,
+        **kwargs,
+    ) -> None:
+        """
+        Piecewise-constant linear LR schedule updated every `decay_steps`.
+
+        The LR factor linearly interpolates from `start_factor` (at step=0)
+        to `end_factor` (at and after step >= stop_steps), but the value only
+        changes at discrete update boundaries (multiples of `decay_steps`).
+
+        Parameters
+        ----------
+        start_lr : float
+            Base learning rate (multiplied by the factor below).
+        stop_steps : int
+            Total number of training steps for this scheduler.
+        decay_steps : int
+            Interval (in steps) between LR updates; e.g., 1k or 10k.
+        start_factor : float
+            Multiplicative factor at step 0.
+        end_factor : float
+            Multiplicative factor at and after step >= stop_steps.
+
+        Examples
+        --------
+        Let k = floor(step / decay_steps).
+        Let U = stop_steps / decay_steps  (can be non-integer).
+        progress = clamp(k / U, 0, 1).
+        factor(step) = start_factor + (end_factor - start_factor) * progress.
+        After step >= stop_steps, factor(step) = end_factor.
+        - If `decay_steps` >= `stop_steps`, it will be replaced by a reasonable
+            default so the schedule still updates multiple times.
+        - This mirrors the spirit of torch.optim.lr_scheduler.LinearLR but with
+            discrete updates every `decay_steps` steps (akin to treating each
+            update as an "epoch").
+        """
+        self.base_lr = float(start_lr)
+        self.start_factor = float(start_factor)
+        self.end_factor = float(end_factor)
+        self.stop_steps = int(stop_steps)
+
+        # Choose a safe decay_steps (avoid zero/oversized intervals)
+        self.decay_steps = int(decay_steps) if int(decay_steps) > 0 else 1
+        default_ds = 100 if self.stop_steps // 10 > 100 else self.stop_steps // 100 + 1
+        if self.decay_steps >= self.stop_steps:
+            self.decay_steps = max(1, int(default_ds))
+
+        # Total number of "update buckets" over the training horizon (float)
+        self.total_updates = self.stop_steps / self.decay_steps
+
+    def value(self, step: int) -> np.float64:
+        """
+        Get the learning rate at the given `step`.
+
+        - Updates occur only at multiples of `decay_steps`.
+        - Saturates at `end_factor` when step >= stop_steps.
+        - Negative steps are treated as 0.
+        """
+        if step <= 0:
+            factor = self.start_factor
+        elif step >= self.stop_steps:
+            factor = self.end_factor
+        else:
+            updates_done = step // self.decay_steps  # integer count of updates so far
+            progress = (
+                updates_done / self.total_updates
+            )  # may be slightly < 1 before stop_steps
+            # Clamp numerical drift into [0, 1]
+            if progress < 0.0:
+                progress = 0.0
+            elif progress > 1.0:
+                progress = 1.0
+
+            factor = (
+                self.start_factor + (self.end_factor - self.start_factor) * progress
+            )
+            # Monotone clamp to never overshoot end_factor due to rounding
+            if self.end_factor < self.start_factor:
+                factor = max(factor, self.end_factor)
+            else:
+                factor = min(factor, self.end_factor)
+
+        return np.float64(self.base_lr * factor)

deepmd/pt/train/training.py

@@ -61,6 +61,7 @@
 from deepmd.pt.utils.learning_rate import (
     LearningRateCosine,
     LearningRateExp,
+    LearningRateLinear,
     LearningRateWSD,
 )
 from deepmd.pt.utils.stat import (
@@ -252,6 +253,8 @@ def get_lr(lr_params):
                 lr_schedule = LearningRateCosine(**lr_params)
             elif lr_type == "wsd":
                 lr_schedule = LearningRateWSD(**lr_params)
+            elif lr_type == "linear":
+                lr_schedule = LearningRateLinear(**lr_params)
             else:
                 raise ValueError(f"Not supported learning rate type '{lr_type}'!")
             return lr_schedule

deepmd/pt/utils/learning_rate.py

@@ -2,11 +2,13 @@
 from deepmd.dpmodel.utils.learning_rate import (
     LearningRateCosine,
     LearningRateExp,
+    LearningRateLinear,
     LearningRateWSD,
 )
 
 __all__ = [
     "LearningRateCosine",
     "LearningRateExp",
+    "LearningRateLinear",
     "LearningRateWSD",
 ]

deepmd/utils/argcheck.py

@@ -3066,6 +3066,17 @@ def learning_rate_wsd():
     return args
 
 
+def learning_rate_linear():
+    doc_start_lr = "The learning rate at the start of the training."
+    args = [
+        Argument("start_lr", float, optional=True, default=1e-3, doc=doc_start_lr),
+        Argument("start_factor", float, optional=True, default=1.0),
+        Argument("end_factor", float, optional=True, default=1e-3),
+        Argument("decay_steps", int, optional=True, default=1000),
+    ]
+    return args
+
+
 def learning_rate_variant_type_args():
     doc_lr = "The type of the learning rate."
 
@@ -3075,6 +3086,7 @@ def learning_rate_variant_type_args():
             Argument("exp", dict, learning_rate_exp()),
             Argument("cosine", dict, learning_rate_cosine()),
             Argument("wsd", dict, learning_rate_wsd()),
+            Argument("linear", dict, learning_rate_linear()),
         ],
         optional=True,
         default_tag="exp",