Add learning_rate to the summary.

PiperOrigin-RevId: 313148142

Add learning_rate to the summary.
PiperOrigin-RevId: 313148142
ee760015 · Maxim Neumann · A. Unique TensorFlower · 38800db3 · ee760015
Commit ee760015 authored May 26, 2020 by Maxim Neumann Committed by A. Unique TensorFlower May 26, 2020
Show whitespace changes
Inline Side-by-side

Showing with 5 additions and 0 deletions

official/nlp/bert/model_training_utils.py official/nlp/bert/model_training_utils.py +5 -0

No files found.
--- a/official/nlp/bert/model_training_utils.py
+++ b/official/nlp/bert/model_training_utils.py
@@ -488,6 +488,11 @@ def run_customized_training_loop(
        summary_writer = tf.summary.create_noop_writer()

      with summary_writer.as_default():
+        if callable(optimizer.learning_rate):
+          tf.summary.scalar(
+              'learning_rate',
+              optimizer.learning_rate(current_step),
+              step=current_step)
        tf.summary.scalar(
            train_loss_metric.name, train_loss, step=current_step)
        for metric in train_metrics + model.metrics: