Cannot compose Augmentations with the `keras_core.Sequential` API

[ariG23498]

I was using the following code for a semantic segmentation pipeline

augment_fn = keras.Sequential(
    [
        resize_fn,
        keras_cv.layers.RandomFlip(),
        keras_cv.layers.RandomRotation(
            factor=ROTATION_FACTOR,
            segmentation_classes=NUM_CLASSES,
        ),
        keras_cv.layers.RandAugment(
            value_range=(0, 1),
            geometric=False,
        ),
    ]
)

The idea was to load a dataset using the tf.data.Dataset API, and then .map the dataset with the augment_fn model.

Doing that gives a NotImplementedError

NotImplementedError : Cannot convert a symbolic tf.Tensor (args_0:0) to a numpy array. This error may indicate that you're trying to pass a Tensor to a NumPy call, which is not supported.

To reproduce the error: Colab Notebook

[ianstenbit]

@ariG23498 thanks for the error report -- are you currently able to workaround by doing something like this:

augmentations = [keras_cv.layers.RandAugment(...), ...]
def augment(inputs):
    for layer in augmentations:
        inputs = layer(inputs)
    return inputs
dataset.map(augment)

I'd also expect that this should still work if you use tf.keras.Sequential.

To be clear -- this is a bug and I'll dig into it when I can. Just trying to establish whether there's a workaround for now.

[ariG23498]

@ianstenbit

Thanks for the workaround! The functional approach definitely helps! I did check it out and it ran perfectly.

[jbischof]

It appears that tf.keras.Sequential does not work, at least if KerasCV is using Keras Core. When trying to create the model I get the error:

TypeError: The added layer must be an instance of class Layer. Received: layer=<RandomFlip name=random_flip_3, built=False> of type <class 'keras_cv.layers.preprocessing.random_flip.RandomFlip'>.

See gist

[mattdangerw]

See gist

This is actually an expected sharp edge. Mixing a tf.keras.Model with a keras_core.layers.Layer is a no go. If you have os.environ['KERAS_BACKEND'] = 'jax', keras_cv.layers.RandomFlip() is going to be a keras_core.layers.Layer. Which tf.keras.Model knows absolutely nothing about.

For now, this is just something we have to hold in our heads. Don't mix tf.keras and keras-core!

This isn't keras-cv specific, the same blow up would happen with tf.keras and keras-core directly. But it does get extra confusing with keras-cv, as keras-cv can load up all it's layers as either tf.keras.layers.Layer or keras_core.layers.Layer depending how it is configured. Confusing for sure! If we can figure out and error/warning that would help out.

Thankfully this is not a forever state.

[jbischof]

Talking with the team we should re-introduce the Augmenter API to reduce boilerplate for our users in this case.

We had provided this functionality earlier but deprecated it in favor of the now broken keras.Model wrapper.