'How can you update LR scheduler with callback in tensorflow?

I am trying to pick up a model after the training and run several more epochs with LR callback.

The problem I am encountering is that the original learning rate, which is CosineDecay, cannot get updated with the function tf.keras.backend.set_value. The model learning rate is shown as below:

>> model.optimizer.learning_rate

<keras.optimizer_v2.learning_rate_schedule.CosineDecay at 0x7fb46b2f6850>

Here is the callback function I am using:

(Modified from: Learning Rate Callback on Step Rather Than Epoch? ).

class UpdateLR_byStep(tf.keras.callbacks.Callback):
    def __init__(self, schedule, steps_per_epoch):
        super(UpdateLR_byStep, self).__init__()
        self.schedule = schedule
        self.epoch = 0
        self.steps_per_epoch = steps_per_epoch
        
    def on_batch_begin(self, batch, logs=None):
        actual_step = (self.epoch*self.steps_per_epoch) + batch
        scheduled_lr = self.schedule(actual_step)
        tf.keras.backend.set_value(self.model.optimizer.lr, scheduled_lr)
        if batch == 0:
            print("--Learning Rate: {:.3g} --".format(scheduled_lr))
        self.scheduled_lr = scheduled_lr
        
    def on_epoch_end(self, epoch, logs=None):
        message = '  LR: {}.'.format(self.scheduled_lr)
        sys.stdout.write(message + "\n")
        self.epoch+=1

I pinned down the problem it is giving me an error because I am trying to update learningRateScheduler class, instead of a number.

For example, if I try to update the learning rate with a floating number 1e-7, it gives me an error:


tf.keras.backend.set_value(model.optimizer.learning_rate, 1e-7)

---------------------------------------------------------------------------

AttributeError                            Traceback (most recent call last)

<ipython-input-96-5701462812c1> in <module>()
      3   )
      4 
----> 5 tf.keras.backend.set_value(model.optimizer.learning_rate, 1e-7)

1 frames

/usr/local/lib/python3.7/dist-packages/keras/backend.py in set_value(x, value)
   3982           (of the same shape).
   3983   """
-> 3984   value = np.asarray(value, dtype=dtype_numpy(x))
   3985   if tf.compat.v1.executing_eagerly_outside_functions():
   3986     x.assign(value)

/usr/local/lib/python3.7/dist-packages/keras/backend.py in dtype_numpy(x)
   1546       numpy.dtype, dtype of `x`.
   1547   """
-> 1548   return tf.as_dtype(x.dtype).as_numpy_dtype
   1549 
   1550 

AttributeError: 'CosineDecay' object has no attribute 'dtype'

Any way to work around this? I don't want to recompile the model.



Sources

This article follows the attribution requirements of Stack Overflow and is licensed under CC BY-SA 3.0.

Source: Stack Overflow

Solution Source