Repository: keras-team/keras-nlp

Issue:
Add a proper learning rate schedule to the BERT example
We should add learning rate schedule matching the [original BERT implementation](https://github.com/google-research/bert/blob/master/optimization.py#L32) to our BERT example in `/examples/bert`.


Discussion:
https://keras.io/api/callbacks/learning_rate_scheduler/
Hey @mattdangerw! I'd like to work on this issue.
Hey @mattdangerw! I just wanted to confirm a couple of things which I'll be doing:
- The file where the issue is focused is `/examples/bert/run_pretraining.py`.
- Adding a `scheduler` function as referenced [here](https://keras.io/api/callbacks/learning_rate_scheduler/) and then use the same to update the `learning_rate` like so: `learning_rate_schedule = scheduler(epoch = FLAGS.epochs, lr = FLAGS.learning_rate)`.
Am I correct?
Thanks for contributing!

Yes this should be changing `/examples/bert/run_pretraining.py`

Re scheduler, sorry I linked the wrong resource there. We need to set the learning rate depending on the number of train steps, not epochs, so we should use the following:
https://www.tensorflow.org/api_docs/python/tf/keras/optimizers/schedules/LearningRateSchedule

I believe we want a linear ramp up during the first 10k training steps, followed by a linear decay. The code I linked shows the logic for original BERT, though it will need to be adapted to work with Keras/TF2.
Thank you for the review!

I think I understand, so currently I'm thinking of implementing the scheduler by increasing the learning rate with a factor of `initial_learning_rate/(current_step + 1)`, and for decaying, would a factor shallower than the increasing factor be fine, till an `end_learning rate` of let's say `0.0001`?

Resolving patch:
diff --git a/examples/bert/run_pretraining.py b/examples/bert/run_pretraining.py
--- a/examples/bert/run_pretraining.py
+++ b/examples/bert/run_pretraining.py
@@ -61,6 +61,19 @@
     "Maximum number of masked LM predictions per sequence.",
 )
 
+flags.DEFINE_integer(
+    "num_warmup_steps",
+    1e4,
+    "The number of warmup steps during which the learning rate will increase "
+    "till a threshold.",
+)
+
+flags.DEFINE_integer(
+    "num_train_steps",
+    1e6,
+    "The total fixed number of steps till which the model will train.",
+)
+
 
 class ClassificationHead(tf.keras.layers.Layer):
     """Pooling head for sentence-level classification tasks.
@@ -292,6 +305,34 @@ def train_step(self, data):
         return {m.name: m.result() for m in self.metrics}
 
 
+class LinearDecayWithWarmup(keras.optimizers.schedules.LearningRateSchedule):
+    """
+    A learning rate schedule with linear warmup and decay.
+
+    This schedule implements a linear warmup for the first `num_warmup_steps`
+    and a linear ramp down until `num_train_steps`.
+    """
+
+    def __init__(self, learning_rate, num_warmup_steps, num_train_steps):
+        self.learning_rate = learning_rate
+        self.warmup_steps = num_warmup_steps
+        self.train_steps = num_train_steps
+
+    def __call__(self, step):
+        peak_lr = tf.cast(self.learning_rate, dtype=tf.float32)
+        warmup = tf.cast(self.warmup_steps, dtype=tf.float32)
+        training = tf.cast(self.train_steps, dtype=tf.float32)
+
+        is_warmup = step < warmup
+
+        # Linear Warmup will be implemented if current step is less than
+        # `num_warmup_steps`.
+        if is_warmup:
+            return peak_lr * (step / warmup)
+        # else Linear Decay will be implemented
+        return max(0.0, peak_lr * (training - step) / (training - warmup))
+
+
 def decode_record(record):
     """Decodes a record to a TensorFlow example."""
     seq_length = FLAGS.max_seq_length
@@ -348,14 +389,22 @@ def main(_):
     model(model.inputs)
     model.summary()
 
+    learning_rate_schedule = LinearDecayWithWarmup(
+        learning_rate=FLAGS.learning_rate,
+        num_warmup_steps=FLAGS.num_warmup_steps,
+        num_train_steps=FLAGS.num_train_steps,
+    )
+
     # Wrap with pretraining heads and call fit.
     pretraining_model = BertPretrainer(model)
     pretraining_model.compile(
-        # TODO(mattdangerw): Add AdamW and a learning rate schedule.
-        optimizer=keras.optimizers.Adam(learning_rate=FLAGS.learning_rate)
+        optimizer=keras.optimizers.Adam(learning_rate=learning_rate_schedule)
     )
     # TODO(mattdangerw): Add TPU strategy support.
-    pretraining_model.fit(dataset, epochs=FLAGS.epochs)
+    steps_per_epoch = FLAGS.num_train_steps // FLAGS.epochs
+    pretraining_model.fit(
+        dataset, epochs=FLAGS.epochs, steps_per_epoch=steps_per_epoch
+    )
 
     print(f"Saving to {FLAGS.saved_model_output}")
     model.save(FLAGS.saved_model_output)


Test patch: