Keras TPU. Compilation failure: Detected unsupported operations

Question

I try to run my keras UNet model using Google Colab TPU and I faced this problem with UpSampling2D. Any solutions or workaround?

Code to run:

import os
import numpy as np
import tensorflow as tf
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import UpSampling2D

model = Sequential()
model.add(UpSampling2D((2, 2), input_shape=(16, 16, 1)))
model.compile(optimizer=tf.train.RMSPropOptimizer(learning_rate=0.01), 
              loss='binary_crossentropy', metrics=['acc'])

TPU_WORKER = 'grpc://' + os.environ['COLAB_TPU_ADDR']
tf.logging.set_verbosity(tf.logging.INFO)


model = tf.contrib.tpu.keras_to_tpu_model(
    model,strategy=tf.contrib.tpu.TPUDistributionStrategy(
        tf.contrib.cluster_resolver.TPUClusterResolver(TPU_WORKER)))

X = np.zeros((1024, 16, 16, 1))
Y = np.zeros((1024, 32, 32, 1))

model.fit(X, Y, batch_size=1024)

Error:

RuntimeError: Compilation failed: Compilation failure: Detected unsupported operations when trying to compile graph cluster_3_5095732716396540171[] on XLA_TPU_JIT: ResizeNearestNeighbor (No registered 'ResizeNearestNeighbor' OpKernel for XLA_TPU_JIT devices compatible with node {{node tpu_140211339657168/up_sampling2d_1/ResizeNearestNeighbor}} = ResizeNearestNeighbor[T=DT_FLOAT, align_corners=false, _device="/device:TPU_REPLICATED_CORE"](infeed-train_1:1, tpu_140211339657168/up_sampling2d_1/mul) . Registered: device='CPU'; T in [DT_DOUBLE] device='CPU'; T in [DT_FLOAT] device='CPU'; T in [DT_BFLOAT16] device='CPU'; T in [DT_HALF] device='CPU'; T in [DT_INT8] device='CPU'; T in [DT_UINT8] device='CPU'; T in [DT_INT16] device='CPU'; T in [DT_UINT16] device='CPU'; T in [DT_INT32] device='CPU'; T in [DT_INT64] ){{node tpu_140211339657168/up_sampling2d_1/ResizeNearestNeighbor}}

Answer 1

From the error it looks like one of the operations in your Tensorflow backend (ResizeNearestNeighbor) graph for Keras is currently not compatible with TPUs. There are a small number of Tensorflow ops that are currently not available for TPUs (Cloud TPU FAQs).

You can check out the current list of TPU compatible Tensorflow ops here. You can also use Tensorboard to see TPU Compatibility Graphs.

As a workaround, you can try combine the TPU compatible Tensorflow ops to replicate the behavior of ResizeNearestNeighbor. In particular, you may be interested in the ResizeBilinear Op, which is TPU compatible.