在colab TPU上运行keras模型

时间:2019-12-11 16:21:45

标签: python keras conv-neural-network google-colaboratory tpu

我目前正在尝试使用Keras和CNN GPU训练卷积神经网络 Google Colab 。 我发现这个article讨论了增加训练模型所需的训练时间的选项。由于目前 GPU 的培训非常慢,因此我尝试实现本文中的方法。我有以下代码:

sgd = optimizers.SGD(lr=0.02)
model.compile(optimizer=sgd,loss='categorical_crossentropy',metrics=['accuracy'])
def create_train_subsets():
    X_train =[]
    y_train = []
    for i in range(80):
      cat = i+1
      path = 'train_set/by_cat/{}'.format(cat)
      for img in os.listdir(path):
        actual_image = Image.open(("train_set/by_cat/{}/{}".format(cat,img)))
        X_train.append(actual_image)
        y_train.append(cat)
    return X_train, y_train

# This address identifies the TPU we'll use when configuring TensorFlow.
x_train, y_train = create_train_subsets()
TPU_WORKER = 'grpc://' + os.environ['COLAB_TPU_ADDR']
tf.logging.set_verbosity(tf.logging.INFO)

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

history = tpu_model.fit(x_train, y_train,
                        epochs=20,
                        batch_size=128 * 8,
                        validation_split=0.2)

tpu_model.save_weights('./tpu_model.h5', overwrite=True)
# tpu_model.evaluate(x_test, y_test, batch_size=128 * 8)

但是此代码返回以下错误:

InvalidArgumentError: No OpKernel was registered to support Op 'ConfigureDistributedTPU' used by node ConfigureDistributedTPU (defined at /usr/local/lib/python3.6/dist-packages/tensorflow_core/python/framework/ops.py:1748) with these attrs: [tpu_embedding_config="", is_global_init=false, embedding_config=""]
Registered devices: [CPU, XLA_CPU]
Registered kernels:
  <no registered kernels>

     [[ConfigureDistributedTPU]]

我在网上进行了广泛的搜索,但似乎找不到任何含义。另外,我对这个过程还不够了解,无法自己弄清错误的确切含义。

因此,有没有人可以帮助我了解问题所在,也许还知道解决该问题的方法。

提前谢谢!

0 个答案:

没有答案