Edit: See end of this question for the Solution

TL;DR: I need to find a way to calculate the label distribution per-batch, and update the learning rate. Is there a way to access the optimizer of the current model to update the learning_rate, per batch?

Below is how to calculate the label distribution. It can be done in the loss function, as by default the loss is calculated batch-wise. Where can this code be executed which also has access to the model's optimizer?

def loss(y_true, y_pred):
    y = math_ops.argmax(y_true, axis=1)
    freqs = tf.gather(lf, y)  # equal to lf[y] if `lf` and `y` were numpy array's
    inv_freqs = math_ops.pow(freqs, -1)
    E = 1 / math_ops.reduce_sum(inv_freqs)  # value to use when updating learning rate

Further Details

In order to implement a learning rate schedule, as described in this paper, I believe I need a way to update the learning rate during training, each batch, by a value calcuated from the label distribution of the true labels in the batch (y_true as it's typically denoted in keras/tensorflow)

where ...

x the output from model

y the corresponding ground truth labels

Β the minibatch of m samples (e.g. 64)

n_y the entire training sample size for ground truth label y

n_y^-1 the inverse label frequency

The portion of the formula I'm focused on is the part between α and Δ_θ

I can achieve this with ease from within a custom loss function, but I do not know how to upadte the learning rate--if you even can--from the loss function.

def loss(y_true, y_pred):
    y = math_ops.argmax(y_true, axis=1)
    freqs = tf.gather(lf, y)  # equal to lf[y] if `lf` and `y` were numpy array's
    inv_freqs = math_ops.pow(freqs, -1)
    E = 1 / math_ops.reduce_sum(inv_freqs)  # value to use when updating learning rate

where ...

lf the sample frequencies for each class. e.g. 2 classes, c0 = 10 examples, c1 = 100 --> lf == [10, 100]

Is there some fancy way I can update the optimizers learning rate, like what can be done from a CallBack?

def on_batch_begin(self, batch, log):
    # note: batch is just an incremented value to indicate batch index
    self.model.optimizer.lr  # learning rate, can be modified from callback

Thanks in advance for any help!

SOLUTION

Huge thank you to @mrk for pushing me in the right direction to solve this!

In order to compute the per-batch label distributions, then use that value to update the optimizer's learning rate, one must ...

1. Create a custom Metric which computes the label distribution, per-batch, and returns the frequency array (by default keras is optimized batch-wise, hence metrics are calcuated each batch).
2. Create a typical learning rate scheduler, by subclassing the keras.callbacks.History class
3. Override the on_batch_end function of the scheduler, the logs dict will ontain all computed metrics for the batch including our custom label distribution metric!

Creating Custom Metric

class LabelDistribution(tf.keras.metrics.Metric):
    """
    Computes the per-batch label distribution (y_true) and stores the array as
    a metric which can be accessed via keras CallBack's

    :param n_class: int - number of distinct output class(es)
    """

    def __init__(self, n_class, name='batch_label_distribution', **kwargs):
        super(LabelDistribution, self).__init__(name=name, **kwargs)
        self.n_class = n_class
        self.label_distribution = self.add_weight(name='ld', initializer='zeros',
                                                  aggregation=VariableAggregation.NONE,
                                                  shape=(self.n_class, ))

    def update_state(self, y_true, y_pred, sample_weight=None):
        y_true = mo.cast(y_true, 'int32')
        y = mo.argmax(y_true, axis=1)
        label_distrib = mo.bincount(mo.cast(y, 'int32'))

        self.label_distribution.assign(mo.cast(label_distrib, 'float32'))

    def result(self):
        return self.label_distribution

    def reset_states(self):
        self.label_distribution.assign([0]*self.n_class)

Create DRW Learning Rate Scheduler

class DRWLearningRateSchedule(keras.callbacks.History):
    """
    Used to implement the Differed Re-weighting strategy from
    [Kaidi Cao, et al. "Learning Imbalanced Datasets with Label-Distribution-Aware Margin Loss." (2019)]
    (https://arxiv.org/abs/1906.07413)

    To be included as a metric to model.compile
    `model.compile(..., metrics=[DRWLearningRateSchedule(.01)])`
    """

    def __init__(self, base_lr, ld_metric='batch_label_distribution'):
        super(DRWLearningRateSchedule, self).__init__()

        self.base_lr = base_lr
        self.ld_metric = ld_metric  # name of the LabelDistribution metric

    def on_batch_end(self, batch, logs=None):
        ld = logs.get(self.ld_metric)  # the per-batch label distribution
        current_lr = self.model.optimizer.lr
        # example below of updating the optimizers learning rate
        K.set_value(self.model.optimizer.lr, current_lr * (1 / math_ops.reduce_sum(ld)))