Repository: OpenNMT/OpenNMT-tf
Pull request #961 (2022-07-28 16:12:34)

Make initial learning rate of InvSqrtDecay configurable


Patch:
diff --git a/opennmt/schedules/lr_schedules.py b/opennmt/schedules/lr_schedules.py
index 7dc824ae..d4e44546 100644
--- a/opennmt/schedules/lr_schedules.py
+++ b/opennmt/schedules/lr_schedules.py
@@ -175,7 +175,9 @@ class InvSqrtDecay(tf.keras.optimizers.schedules.LearningRateSchedule):
 
     .. math::
 
-        \text{schedule}(\text{step}) = \text{learning_rate}
+        \text{schedule}(\text{step}) = \text{init_lr}
+                                       +
+                                       (\text{lr} - \text{init_lr})
                                        \times
                                        \frac{\text{step}}{\text{warmup_steps}}
 
@@ -183,7 +185,7 @@ class InvSqrtDecay(tf.keras.optimizers.schedules.LearningRateSchedule):
 
     .. math::
 
-        \text{schedule}(\text{step}) = \text{learning_rate}
+        \text{schedule}(\text{step}) = \text{lr}
                                        \times
                                        \sqrt{\frac{\text{warmup_steps}}{\text{step}}}
 
@@ -191,22 +193,31 @@ class InvSqrtDecay(tf.keras.optimizers.schedules.LearningRateSchedule):
       - :class:`opennmt.schedules.RsqrtDecay`
     """
 
-    def __init__(self, learning_rate, warmup_steps):
+    def __init__(self, learning_rate, warmup_steps, initial_learning_rate=0):
         """Initializes the decay function.
 
         Args:
           learning_rate: The base learning rate.
           warmup_steps: The number of warmup steps.
+          initial_learning_rate: Initial learning rate during warmup.
         """
-        self.learning_rate = tf.cast(learning_rate, tf.float32)
+        self.lr = tf.cast(learning_rate, tf.float32)
+        self.init_lr = tf.cast(initial_learning_rate, tf.float32)
         self.warmup_steps = tf.cast(warmup_steps, tf.float32)
 
     def __call__(self, step):
         step = tf.cast(step + 1, tf.float32)
-        return self.learning_rate * tf.cond(
+
+        def _warmup():
+            return self.init_lr + (self.lr - self.init_lr) * (step / self.warmup_steps)
+
+        def _after_warmup():
+            return self.lr * tf.math.sqrt(self.warmup_steps / step)
+
+        return tf.cond(
             step <= self.warmup_steps,
-            true_fn=lambda: step / self.warmup_steps,
-            false_fn=lambda: tf.math.sqrt(self.warmup_steps / step),
+            true_fn=_warmup,
+            false_fn=_after_warmup,
         )
 
 


Test patch:
diff --git a/opennmt/tests/lr_schedules_test.py b/opennmt/tests/lr_schedules_test.py
index 4c995cbd..7428c397 100644
--- a/opennmt/tests/lr_schedules_test.py
+++ b/opennmt/tests/lr_schedules_test.py
@@ -1,5 +1,7 @@
 import tensorflow as tf
 
+from parameterized import parameterized
+
 from opennmt.schedules import lr_schedules
 
 
@@ -69,11 +71,14 @@ class LRSchedulesTest(tf.test.TestCase):
     def testRsqrtDecay(self):
         self._testNoError(lr_schedules.RsqrtDecay(2.0, 4000))
 
-    def testInvSqrtDecay(self):
+    @parameterized.expand([(0,), (1e-07,)])
+    def testInvSqrtDecay(self, initial_learning_rate):
         learning_rate = 0.0002
         warmup_steps = 4000
-        schedule = lr_schedules.InvSqrtDecay(learning_rate, warmup_steps)
-        self.assertNotEqual(schedule(0), 0)
+        schedule = lr_schedules.InvSqrtDecay(
+            learning_rate, warmup_steps, initial_learning_rate=initial_learning_rate
+        )
+        self.assertNotEqual(schedule(0), initial_learning_rate)
         self.assertEqual(schedule(warmup_steps - 1), learning_rate)
 
     def testCosineAnnealing(self):