Vijay Mariappan
Vijay Mariappan

Reputation: 17201

Data augmentation on GPU

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.

Upvotes: 3

Views: 2419

Answers (2)

aliassaila
aliassaila

Reputation: 83

See the Tensorflow guide that discusses preprocessing data before the model or inside the model. By including preprocessing inside the model, the GPU is leveraged instead of the CPU, it makes the model portable, and it helps reduce the training/serving skew. The guide also has multiple recipes to get you started too. It doesn't explicitly state this works for a TPU but it can be tried.

Upvotes: 2

user11530462
user11530462

Reputation:

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.

Upvotes: 2

Related Questions