How to change a learning rate for Adam in TF2?

Question

How to change the learning rate of Adam optimizer, while learning is progressing in TF2? There some answers floating around, but applicable to TF1, e.g. using feed_dict.

Answer 1

If you are using custom training loop (instead of keras.fit()), you can simply do:

new_learning_rate = 0.01 
my_optimizer.lr.assign(new_learning_rate)

Answer 2

You can read and assign the learning rate via a callback. So you can use something like this:

class LearningRateReducerCb(tf.keras.callbacks.Callback):

  def on_epoch_end(self, epoch, logs={}):
    old_lr = self.model.optimizer.lr.read_value()
    new_lr = old_lr * 0.99
    print("\nEpoch: {}. Reducing Learning Rate from {} to {}".format(epoch, old_lr, new_lr))
    self.model.optimizer.lr.assign(new_lr)

Which, for example, using the MNIST demo can be applied like this:

mnist = tf.keras.datasets.mnist

(x_train, y_train), (x_test, y_test) = mnist.load_data()
x_train, x_test = x_train / 255.0, x_test / 255.0

model = tf.keras.models.Sequential([
  tf.keras.layers.Flatten(input_shape=(28, 28)),
  tf.keras.layers.Dense(128, activation='relu'),
  tf.keras.layers.Dropout(0.2),
  tf.keras.layers.Dense(10, activation='softmax')
])

model.compile(optimizer='adam',
              loss='sparse_categorical_crossentropy',
              metrics=['accuracy'])

model.fit(x_train, y_train, callbacks=[LearningRateReducerCb()], epochs=5)

model.evaluate(x_test, y_test)

giving output like this:

Train on 60000 samples
Epoch 1/5
59744/60000 [============================>.] - ETA: 0s - loss: 0.2969 - accuracy: 0.9151
Epoch: 0. Reducing Learning Rate from 0.0010000000474974513 to 0.0009900000877678394
60000/60000 [==============================] - 6s 92us/sample - loss: 0.2965 - accuracy: 0.9152
Epoch 2/5
59488/60000 [============================>.] - ETA: 0s - loss: 0.1421 - accuracy: 0.9585
Epoch: 1. Reducing Learning Rate from 0.0009900000877678394 to 0.000980100128799677
60000/60000 [==============================] - 5s 91us/sample - loss: 0.1420 - accuracy: 0.9586
Epoch 3/5
59968/60000 [============================>.] - ETA: 0s - loss: 0.1056 - accuracy: 0.9684
Epoch: 2. Reducing Learning Rate from 0.000980100128799677 to 0.0009702991228550673
60000/60000 [==============================] - 5s 91us/sample - loss: 0.1056 - accuracy: 0.9684
Epoch 4/5
59520/60000 [============================>.] - ETA: 0s - loss: 0.0856 - accuracy: 0.9734
Epoch: 3. Reducing Learning Rate from 0.0009702991228550673 to 0.0009605961386114359
60000/60000 [==============================] - 5s 89us/sample - loss: 0.0857 - accuracy: 0.9733
Epoch 5/5
59712/60000 [============================>.] - ETA: 0s - loss: 0.0734 - accuracy: 0.9772
Epoch: 4. Reducing Learning Rate from 0.0009605961386114359 to 0.0009509901865385473
60000/60000 [==============================] - 5s 87us/sample - loss: 0.0733 - accuracy: 0.9772
10000/10000 [==============================] - 0s 43us/sample - loss: 0.0768 - accuracy: 0.9762
[0.07680597708942369, 0.9762]