Keras TPU。编译失败:检测到不支持的操作

时间:2018-10-20 13:28:59

标签: tensorflow keras google-colaboratory google-cloud-tpu

我尝试使用Google Colab TPU运行我的keras UNet模型,并且遇到UpSampling2D这个问题。有解决方案或解决方法吗?

要运行的代码:

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)

错误:

  

RuntimeError:编译失败:编译失败:检测到   尝试编译图形时不支持的操作   XLA_TPU_JIT上的cluster_3_5095732716396540171 []:ResizeNearestNeighbor   (没有为XLA_TPU_JIT注册“ ResizeNearestNeighbor” OpKernel   与节点{{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)         。已注册:device ='CPU'; [DT_DOUBLE]中的T         设备='CPU'; [DT_FLOAT]中的T         设备='CPU'; [DT_BFLOAT16]中的T         设备='CPU'; [DT_HALF]中的T         设备='CPU'; [DT_INT8]中的T         设备='CPU'; [DT_UINT8]中的T         设备='CPU'; [DT_INT16]中的T         设备='CPU'; [DT_UINT16]中的T         设备='CPU'; [DT_INT32]中的T         设备='CPU'; [DT_INT64]中的T       ){{node tpu_140211339657168 / up_sampling2d_1 / ResizeNearestNeighbor}}

1 个答案:

答案 0 :(得分:1)

从该错误看来,您的Tensorflow后端(ResizeNearestNeighbor)图中的Keras操作之一当前与TPU不兼容。目前只有少量的Tensorflow操作无法用于TPU(Cloud TPU FAQs)。

您可以查看与TPU兼容的Tensorflow ops here的当前列表。您还可以使用Tensorboard查看TPU Compatibility Graphs

作为一种解决方法,您可以尝试结合使用TPU兼容的Tensorflow操作来复制ResizeNearestNeighbor的行为。特别是,您可能对与TPU兼容的ResizeBilinear Op感兴趣。