How can I create a Keras Learning Rate Schedule that updates based upon batches rather than epochs

I'm working with Keras, and trying to create a Learning Rate Scheduler that schedules on the basis of number of batches processed, instead of number of epochs. To do this, I've inserted the scheduling code into the get_updates method of my `Optimizer'. For the most part, I've tried to use regular Python variables for values that remain constant during a given training run and computational graph nodes only for parameters that actually vary.

My 2 Questions are:

Does the code below look like it should behave properly as a Learning Rate Scheduler, if placed within the get_updates method of a Keras Optimizer.
How could one embed this code in a Class similar to LearningRateScheduler, but which scheduled based upon number of batches, rather than number of epochs?

    #Copying graph node that stores original value of learning rate
    lr = self.lr 

    # Checking whether learning rate schedule is to be used
    if self.initial_lr_decay > 0:
        # this decay mimics exponential decay from 
        # tensorflow/python/keras/optimizer_v2/exponential_decay 

        # Get value of current number of processed batches from graph node
        # and convert to numeric value for use in K.pow()
        curr_batch = float(K.get_value(self.iterations))

        # Create graph node containing lr decay factor
        # Note: self.lr_decay_steps is a number, not a node
        #       self.lr_decay is a node, not a number
        decay_factor =  K.pow(self.lr_decay, (curr_batch / self.lr_decay_steps)) 

        # Reassign lr to graph node formed by
        # product of graph node containing decay factor
        # and graph node containing original learning rate.
        lr = lr * decay_factor

        # Get product of two numbers to calculate number of batches processed
        # in warmup period
        num_warmup_batches = self.steps_per_epoch_num * self.warmup_epochs

        # Make comparisons between numbers to determine if we're in warmup period
        if (self.warmup_epochs > 0) and (curr_batch < num_warmup_batches):

            # Create node with value of learning rate by multiplying a number
            # by a node, and then dividing by a number
            lr = (self.initial_lr  *
                  K.cast(self.iterations, K.floatx()) / curr_batch)

Solution

Easier than messing with Keras source code (it's possible, but it's complex and sensible), you could use a callback.

from keras.callbacks import LambdaCallback

total_batches = 0
def what_to_do_when_batch_ends(batch, logs):
   total_batches += 1 #or use the "batch" variable,
                      #which is the batch index of the last finished batch

   #change learning rate at will
   if your_condition == True:
       keras.backend.set_value(model.optimizer.lr, newLrValueAsPythonFloat)

When training, use the callback:

lrUpdater = LambdaCallback(on_batch_end = what_to_do_when_batch_ends)
model.fit(........, callbacks = [lrUpdater, ...other callbacks...])