BUG: `RandomContrast` slows down training 6-fold

Augmentation will obviously slow down training, but it shouldn’t be a 6-fold slowdown. This happens with the RandomContrast layer, which makes the training per epoch grow from ~100s to ~600s.

I’d share a Colab notebook but there seems to be an issue with KerasCV on Colab on importing, so here are the steps to reproduce:

import keras_cv
print(f'KerasCV version {keras_cv.__version__}')
import tensorflow as tf
print(f'TF version {tf.__version__}')
from tensorflow import keras
print(f'Keras version {keras.__version__}')

(x_train, y_train), (x_test, y_test) = tf.keras.datasets.cifar10.load_data()

aug = keras.Sequential([
    keras_cv.layers.RandomContrast(factor=0.2)
])

model = keras.Sequential([
    keras.layers.InputLayer(input_shape=(32, 32, 3)),
    aug,
    keras.applications.EfficientNetV2B0(weights=None, include_top=False),
    keras.layers.GlobalAveragePooling2D(),
    keras.layers.Dense(10, activation='softmax')
    
])
model.compile(optimizer='adam', loss='sparse_categorical_crossentropy')

model.fit(x_train, y_train)

model = keras.Sequential([
    keras.layers.InputLayer(input_shape=(32, 32, 3)),
    keras.applications.EfficientNetV2B0(weights=None, include_top=False),
    keras.layers.GlobalAveragePooling2D(),
    keras.layers.Dense(10, activation='softmax')
    
])
model.compile(optimizer='adam', loss='sparse_categorical_crossentropy')

model.fit(x_train, y_train)

[bhack]

It is contrast_factor in this example but it is just a case of a more general issue. Just in April the list in Keras-CV was already quite long: https://github.com/tensorflow/tensorflow/issues/55639#issuecomment-1112227581

And this will happen every time you will want to randomize an arg, within the batch, of an Op where that arg is scalar by the TF API design.

Then we could tell that it is a TF issue and we don’t care but who have really the resources to change all these API/ops/kernels in TF for your new design needs?

At the same time Keras is no more a multi backed library, users are impacted directly by the performance issues we have with this design and they cannot switch to an alternative “backend”. So the separation logic between Keras issues and TF issues in many cases doesn’t make sense by an user point of view.

So a limit of your within the batch randomization policy it is going to impact directly the Keras-cv user base and the claim that it is a TF (team?) issue it does not solve anyone’s problems.

At the same time we don’t have (public?) experimental data about the accuracy/epochs gain related to the choice to randomize an augmentation over each batch element that it is the design choice that created all these performance overhead considering the TF API design we currently have on many Ops.

[bhack]

@atlanticstarr1 We have started a thread at https://github.com/tensorflow/tensorflow/issues/55639#issuecomment-1233073292

You could try to ping there.