tensorflow - Data augmentation on GPU

Question

As tf.data augmentations are executed only on CPUs. I need a way to run certain augmentations on the TPU for an audio project.
For example,

CPU: tf.recs read -> audio crop -> noise addition.
TPU: spectogram -> Mixup Augmentation.

Most augmentations can be done as a Keras Layer on top of the model, but MixUp requires both changes in input as well as label.

Is there a way to do it using tf keras APIs.

And if there is any way we can transfer part of tf.data to run on TPU that will also be helpful.

Answer 1

As you have rightly mentioned and as per the Tensorflow documentation also the preprocessing of tf.data is done on CPU only.
However, you can do some workaround to preprocess your tf.data using TPU/GPU by directly using transformation function in your model with something like below code.

input = tf.keras.layers.Input((512,512,3))
x = tf.keras.layers.Lambda(transform)(input) 

You can follow this Kaggle post for detailed discussion on this topic.