Keras RandAugment incorrect behavior in tf graph mode

In tensorflow graph mode, RandAugment will repeatedly use the same augmentations, which were sampled during graph tracing.

It relies on shuffling a python list using random.shuffle, which only works during eager mode execution. In graph mode the operations are sampled then compiled, but the sampling process itself isn't compiled so the same operations are used repeatedly.

https://github.com/keras-team/keras/blob/44a655bdb28037046ab279a49d4cd679fea7ca50/keras/src/layers/preprocessing/image_preprocessing/rand_augment.py#L173

If I add a tf.print statement to this code:

random.shuffle(self._AUGMENT_LAYERS)
for layer_name in self._AUGMENT_LAYERS[: self.num_ops]:
    tf.print(layer_name, tf.executing_eagerly())  # <----
    augmentation_layer = getattr(self, layer_name)
    transformation[layer_name] = (
        augmentation_layer.get_random_transformation(
            data,
            training=training,
            seed=self._get_seed_generator(self.backend._backend),
        )
    )

then run this test:

def test_graph_issue(self):
    input_data = np.random.random((10, 8, 8, 3))
    layer = layers.RandAugment()
    ds = tf_data.Dataset.from_tensor_slices(input_data).batch(2).map(layer)
    print()
    for output in ds:
        output.numpy()

i get this output

equalization False
random_posterization False
equalization False
random_posterization False
equalization False
random_posterization False
equalization False
random_posterization False
equalization False
random_posterization False

Comment From: gregstarr

to clarify, the correct behavior would do random augmentations each time, so the printout would not repeat the same 2 but rather all the augmentations uniformly.

Comment From: fchollet

You're right, we should be using cond operations here, conditioned on variables sampled via keras.random ops. That way the graph will include the conditional branches.

@shashaka what do you think?

Comment From: shashaka

@gregstarr Thank you for reporting the issue. I agree that it should be fixed. As @fchollet mentioned, using a conditional operation seems like a good solution. If you have some time, would you be able to create a PR to address this?

Comment From: gregstarr

yes I will give it a shot

Comment From: gregstarr

can you offer some advice on this PR: #21185 ?

Comment From: sonali-kumari1

Hi @gregstarr - The fix for this issue was addressed in PR #21499 . Please feel free to close this issue if you don't have any further concerns. Thanks!

Comment From: github-actions[bot]

This issue is stale because it has been open for 14 days with no activity. It will be closed if no further activity occurs. Thank you.

Comment From: github-actions[bot]

This issue was closed because it has been inactive for 28 days. Please reopen if you'd like to work on this further.