Colab+TPU not supporting TF 2.3.0 tf.keras.layers.experimental.preprocessing


I was updating my model using TF 2.3.0 on Colab+TPU based on, specifically following the Data augmentation and Transfer learning from pre-trained weights paragraphs.

When I launch I got this error:

InvalidArgumentError: 9 root error(s) found.
  (0) Invalid argument: {{function_node __inference_train_function_372657}} Compilation failure: Detected unsupported operations when trying to compile graph cluster_train_function_12053586239504196919[] on XLA_TPU_JIT: ImageProjectiveTransformV2 (No registered 'ImageProjectiveTransformV2' OpKernel for XLA_TPU_JIT devices compatible with node {{node EfficientNet/img_augmentation/random_rotation_2/transform/ImageProjectiveTransformV2}}){{node EfficientNet/img_augmentation/random_rotation_2/transform/ImageProjectiveTransformV2}}
    TPU compilation failed

I suppose the TPU still does not support tf.keras.layers.experimental.preprocessing because in the list of available TPU operations there is not the preprocessing option. Am I right?

There are multiple Benefits of doing preprocessing inside the model at inference time.

Where could I find a possible implementation date?




A possible workaround is to incorporate the layers into the input pipeline. It’s a bit of a hack, but I’ve tested it briefly and it seems to work on a TPU. For example, if you are using the API, you can create a layer object and then call it within to apply the augmentation to the pipeline:

# dummy data
images = tf.random.uniform((10, 224, 224, 1))
labels = tf.zeros((10, 1))
ds =, labels))
ds = ds.batch(10)

# now incorporate the augmentation 'layer' into the pipeline
augmentor = tf.keras.layers.experimental.preprocessing.RandomRotation((-0.1, 0.1))
# augment the images, pass the labels through untouched
ds = x, y: (, y))

# assume we've compiled a model elsewhere

This doesn’t compile the augmentation layers into the model as originally intended, but it should allow you to augment your training data without requiring a third party plugin. I intend to use this as a workaround until the issue is officially resolved.

Answered By – knuckles

This Answer collected from stackoverflow, is licensed under cc by-sa 2.5 , cc by-sa 3.0 and cc by-sa 4.0

Leave a Reply

(*) Required, Your email will not be published