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}}