Repository: OpenNMT/OpenNMT-tf

Issue:
RsqrtDecay fails when decay_params scale is configured
With following configuration:
```
 decay_type: RsqrtDecay
 decay_params:
   warmup_steps: 10000
   scale: 4
```
Training fails to start with an error:
```
Traceback (most recent call last):
  File "opennmt/bin/main.py", line 350, in <module>
    main()
  File "opennmt/bin/main.py", line 308, in main
    runner.train(
  File "/ai/onmt/onmt-2240/OpenNMT-tf/opennmt/runner.py", line 207, in train
    optimizer = model.get_optimizer()
  File "/ai/onmt/onmt-2240/OpenNMT-tf/opennmt/models/model.py", line 338, in get_optimizer
    learning_rate = schedules.make_learning_rate_schedule(
  File "/ai/onmt/onmt-2240/OpenNMT-tf/opennmt/schedules/lr_schedules.py", line 75, in make_learning_rate_schedule
    schedule = schedule_class(initial_learning_rate, **schedule_params)
TypeError: __init__() got multiple values for argument 'scale'
```
A small change to RsqrtDecay function will fix the issue. It seems that 1st parameter to __init__ is not scale but lr, currently:
```
class RsqrtDecay(tf.keras.optimizers.schedules.LearningRateSchedule):
    """Decay based on the reciprocal of the step square root."""

    def __init__(self ,scale, warmup_steps):
```
If changed to will fix the issue:
```
class RsqrtDecay(tf.keras.optimizers.schedules.LearningRateSchedule):
    """Decay based on the reciprocal of the step square root."""

    def __init__(self, lr ,scale, warmup_steps):
```




Discussion:
In the current implementation, the first argument of learning rate schedules is always the `learning_rate` value defined in the configuration. Some schedules will use this value as the initial learning rate, other schedules will use the value as a scale constant. We could probably improve this behavior.

For now, you could just remove the `scale` parameter and set `learning_rate` instead:

```yaml
params:
  learning_rate: 4
  decay_type: RsqrtDecay
  decay_params:
    warmup_steps: 10000
```

Resolving patch:
diff --git a/opennmt/models/model.py b/opennmt/models/model.py
--- a/opennmt/models/model.py
+++ b/opennmt/models/model.py
@@ -330,12 +330,14 @@ def get_optimizer(self):
         optimizer_name = params.get("optimizer")
         if optimizer_name is None:
             return None
-        learning_rate = tf.constant(params["learning_rate"], dtype=tf.float32)
-        if params.get("decay_type") is not None:
+        schedule_type = params.get("decay_type")
+        if schedule_type is None:
+            learning_rate = tf.constant(params["learning_rate"], dtype=tf.float32)
+        else:
             schedule_params = params.get("decay_params", {})
             learning_rate = schedules.make_learning_rate_schedule(
-                learning_rate,
-                params["decay_type"],
+                params.get("learning_rate"),
+                schedule_type,
                 schedule_params=schedule_params,
                 schedule_step_duration=params.get("decay_step_duration", 1),
                 start_step=params.get("start_decay_steps", 0),
diff --git a/opennmt/schedules/lr_schedules.py b/opennmt/schedules/lr_schedules.py
--- a/opennmt/schedules/lr_schedules.py
+++ b/opennmt/schedules/lr_schedules.py
@@ -1,5 +1,7 @@
 """Define learning rate decay functions."""
 
+import inspect
+
 import numpy as np
 import tensorflow as tf
 
@@ -45,10 +47,10 @@ def make_learning_rate_schedule(
     """Creates the learning rate schedule.
 
     Args:
-      initial_learning_rate: The initial learning rate value or scale.
+      initial_learning_rate: The initial learning rate value. This can be
+        ``None`` if the learning rate is fully defined by the schedule.
       schedule_type: The type of learning rate schedule. A class name from
-        ``tf.keras.optimizers.schedules``
-        or :mod:`opennmt.schedules` as a string.
+        ``tf.keras.optimizers.schedules`` or :mod:`opennmt.schedules` as a string.
       schedule_params: Additional parameters passed to the schedule constructor.
       schedule_step_duration: The number of training steps that make 1 schedule step.
       start_step: Start the schedule after this many steps.
@@ -67,7 +69,10 @@ def make_learning_rate_schedule(
     if schedule_params is None:
         schedule_params = {}
     schedule_class = get_lr_schedule_class(schedule_type)
-    schedule = schedule_class(initial_learning_rate, **schedule_params)
+    first_arg = inspect.getfullargspec(schedule_class)[0][1]
+    if first_arg not in schedule_params:
+        schedule_params[first_arg] = initial_learning_rate
+    schedule = schedule_class(**schedule_params)
     schedule = ScheduleWrapper(
         schedule,
         step_start=start_step,


Test patch:
diff --git a/opennmt/tests/lr_schedules_test.py b/opennmt/tests/lr_schedules_test.py
--- a/opennmt/tests/lr_schedules_test.py
+++ b/opennmt/tests/lr_schedules_test.py
@@ -31,10 +31,18 @@ def testMakeSchedule(self):
         self.assertIsInstance(
             wrapper.schedule, tf.keras.optimizers.schedules.ExponentialDecay
         )
+
         wrapper = lr_schedules.make_learning_rate_schedule(
             2.0, "NoamDecay", dict(model_dim=512, warmup_steps=4000)
         )
         self.assertIsInstance(wrapper.schedule, lr_schedules.NoamDecay)
+        self.assertEqual(wrapper.schedule.scale, 2)
+
+        wrapper = lr_schedules.make_learning_rate_schedule(
+            None, "NoamDecay", dict(scale=2, model_dim=512, warmup_steps=4000)
+        )
+        self.assertEqual(wrapper.schedule.scale, 2)
+
         with self.assertRaises(ValueError):
             lr_schedules.make_learning_rate_schedule(2.0, "InvalidScheduleName")