diff --git a/deepmd/dpmodel/utils/learning_rate.py b/deepmd/dpmodel/utils/learning_rate.py index d26432528c..05c6c9121e 100644 --- a/deepmd/dpmodel/utils/learning_rate.py +++ b/deepmd/dpmodel/utils/learning_rate.py @@ -692,3 +692,30 @@ def _decay_value(self, step: int | Array) -> Array: # Clip to min_lr for steps beyond decay_num_steps step_lr = xp.where(step >= self.decay_num_steps, min_lr, step_lr) return step_lr + + +def make_learning_rate_schedule( + lr_params: dict[str, Any], + num_steps: int, +) -> BaseLR: + """Build a registered learning-rate schedule for a training run. + + The input schema selects schedules through ``learning_rate.type``. Keep + backend trainers on the shared :class:`BaseLR` registry so every backend + accepts the same registered variants without mutating its input config. + + Parameters + ---------- + lr_params : dict[str, Any] + Learning-rate configuration, including the optional ``type`` key. + num_steps : int + Total number of training steps used to parameterize the schedule. + + Returns + ------- + BaseLR + The schedule selected by ``lr_params["type"]`` (``exp`` by default). + """ + params = dict(lr_params) + params["num_steps"] = num_steps + return BaseLR(**params) diff --git a/deepmd/jax/train/trainer.py b/deepmd/jax/train/trainer.py index c19267250e..41c0628e69 100644 --- a/deepmd/jax/train/trainer.py +++ b/deepmd/jax/train/trainer.py @@ -54,7 +54,8 @@ resolve_best_checkpoint_dir, ) from deepmd.dpmodel.utils.learning_rate import ( - LearningRateExp, + BaseLR, + make_learning_rate_schedule, ) from deepmd.dpmodel.utils.multi_task import ( apply_shared_links, @@ -277,11 +278,9 @@ def _deserialize_models(model_data: dict[str, Any]) -> dict[str, BaseModel]: } return {DEFAULT_TASK_KEY: BaseModel.deserialize(model_data["model"])} - def _get_lr_and_coef(self, lr_param: dict[str, Any]) -> LearningRateExp: - lr_type = lr_param.get("type", "exp") - if lr_type == "exp": - return LearningRateExp(**lr_param, num_steps=self.num_steps) - raise RuntimeError("unknown learning_rate type " + lr_type) + def _get_lr_and_coef(self, lr_param: dict[str, Any]) -> BaseLR: + """Construct the schema-selected shared learning-rate schedule.""" + return make_learning_rate_schedule(lr_param, self.num_steps) def _build_losses( self, diff --git a/deepmd/pt_expt/train/training.py b/deepmd/pt_expt/train/training.py index f8dc2d9a1f..e2fce14bb2 100644 --- a/deepmd/pt_expt/train/training.py +++ b/deepmd/pt_expt/train/training.py @@ -40,7 +40,7 @@ split_batch, ) from deepmd.dpmodel.utils.learning_rate import ( - LearningRateExp, + make_learning_rate_schedule, ) from deepmd.pt.train.utils import ( resolve_best_checkpoint_dir, @@ -1476,9 +1476,9 @@ def _make_sample( self.model_prob = None # Learning rate ------------------------------------------------------- - lr_params = config["learning_rate"].copy() - lr_params["num_steps"] = self.num_steps - self.lr_schedule = LearningRateExp(**lr_params) + self.lr_schedule = make_learning_rate_schedule( + config["learning_rate"], self.num_steps + ) # Gradient clipping self.gradient_max_norm = training_params.get("gradient_max_norm", 0.0) diff --git a/deepmd/tf2/train/trainer.py b/deepmd/tf2/train/trainer.py index 8cfa12cbda..1c1205e1ca 100644 --- a/deepmd/tf2/train/trainer.py +++ b/deepmd/tf2/train/trainer.py @@ -47,7 +47,7 @@ split_batch, ) from deepmd.dpmodel.utils.learning_rate import ( - LearningRateExp, + make_learning_rate_schedule, ) from deepmd.dpmodel.utils.training_utils import ( resolve_model_prob, @@ -385,9 +385,9 @@ def sample( resume=init_model is not None or restart_model is not None ) - lr_params = dict(config["learning_rate"]) - lr_params["num_steps"] = self.num_steps - self.lr_schedule = LearningRateExp(**lr_params) + self.lr_schedule = make_learning_rate_schedule( + config["learning_rate"], self.num_steps + ) self.optimizer = self._build_optimizer(config.get("optimizer", {})) self.model_container = _TaskModelContainer(self.models) self.step = tf.Variable(0, dtype=tf.int64, trainable=False, name="step") diff --git a/source/tests/universal/dpmodel/utils/test_learning_rate.py b/source/tests/universal/dpmodel/utils/test_learning_rate.py index 8406f7f352..19551d2165 100644 --- a/source/tests/universal/dpmodel/utils/test_learning_rate.py +++ b/source/tests/universal/dpmodel/utils/test_learning_rate.py @@ -7,9 +7,11 @@ to_numpy_array, ) from deepmd.dpmodel.utils.learning_rate import ( + BaseLR, LearningRateCosine, LearningRateExp, LearningRateWSD, + make_learning_rate_schedule, ) @@ -369,6 +371,29 @@ def test_array_input_wsd_cosine(self) -> None: np.testing.assert_allclose(lrs[3], 1e-5, rtol=1e-10) +class TestMakeLearningRateSchedule(unittest.TestCase): + """Test the shared factory used by backend trainers.""" + + def test_dispatches_all_schema_variants_without_mutating_config(self) -> None: + """Select exp, cosine, and WSD schedules through ``type``.""" + schedule_types: list[tuple[str, type[BaseLR]]] = [ + ("exp", LearningRateExp), + ("cosine", LearningRateCosine), + ("wsd", LearningRateWSD), + ] + for schedule_type, expected_class in schedule_types: + with self.subTest(schedule_type=schedule_type): + params = { + "type": schedule_type, + "start_lr": 1e-3, + "stop_lr": 1e-5, + } + schedule = make_learning_rate_schedule(params, num_steps=100) + + self.assertIsInstance(schedule, expected_class) + self.assertNotIn("num_steps", params) + + class TestLearningRateBeyondStopSteps(unittest.TestCase): """Test learning rate behavior beyond num_steps."""