manifold_mixup_loss_in_keras.py

def manifold_mixup_loss(alpha=1.):
    def _manifold_mixup_loss(y_true, y_pred):
        '''
        y_true: (batch_size, onehot_label_size)
        y_pred: (batch_size, output_unit_size)
        '''

        dist = tf.compat.v1.distributions.Beta(alpha, alpha)
        beta = dist.sample()

        mixuped_y_true = (1. - beta) * tf.reverse(y_true, [0]) + beta * y_true
        mixuped_y_pred = (1. - beta) * tf.reverse(y_pred, [0]) + beta * y_pred

        mixuped_y_pred = tf.nn.softmax(mixuped_y_pred)

        return tf.keras.losses.categorical_crossentropy(mixuped_y_true, mixuped_y_pred)
    return _manifold_mixup_loss