如何在Google Colab中使用TensorFlow 2.0将tf.Keras模型转换为TPU？

Question

由于TF 2.0没有tf.contrib层，如何转换我的模型以在TPU上运行训练，而无权访问tf.contrib.tpu.keras_to_tpu_model()？

我尝试寻找代码，但所有代码都在TensorFlow 1.x上运行

我的数据在中.npy，我有一个简单的模型，我仅在使用它model.compile()并对其model.fit()进行训练，但看起来该模型在CPU上运行（每纪元需要30分钟，而在GPU上则是2分钟）。

Answer 1

使用TensorFlow 2.0将Keras模型转换为TPU（2019年11月更新）

在TensorFlow 2.0中，将Keras Model与Google Cloud TPU一起使用非常容易，不再需要“转换”。

我们需要做的只是指定一个Distributed Strategy使TensorFlow为我们完成所有繁重的工作。

def create_model():
    model = tf.keras.models.Sequential()

    model.add(tf.keras.layers.Conv2D(128, (3, 3), input_shape=x_train.shape[1:]))
    model.add(tf.keras.layers.MaxPooling2D(pool_size=(2, 2), strides=(2,2)))
    model.add(tf.keras.layers.Activation('elu'))

    model.add(tf.keras.layers.Flatten())
    model.add(tf.keras.layers.Dense(10))
    model.add(tf.keras.layers.Activation('softmax'))

    return model

resolver = tf.distribute.cluster_resolver.TPUClusterResolver(tpu='grpc://' + os.environ['COLAB_TPU_ADDR'])
tf.config.experimental_connect_to_host(resolver.master())
tf.tpu.experimental.initialize_tpu_system(resolver)
strategy = tf.distribute.experimental.TPUStrategy(resolver)

with strategy.scope():
    model = create_model()
    model.compile(
        optimizer=tf.keras.optimizers.Adam(learning_rate=1e-3),
        loss=tf.keras.losses.sparse_categorical_crossentropy,
        metrics=[tf.keras.metrics.sparse_categorical_accuracy])

我们创建一个Distribute Resolver，然后为Resolver创建Strategy，然后使用strategy.scope()，创建Keras模型，然后完成。

在https://colab.research.google.com/github/huan/tensorflow-handbook-tpu/blob/master/tensorflow-handbook-tpu-example上了解有关如何使用TPU创建Keras模型的更多信息。 ipynb。

TensorFlow官方分布式培训文件：https : //www.tensorflow.org/guide/distributed_training

但是，请注意，Colab需要解决一些环境问题，因此它可能尚无法在Colab中运行。

Colab TPU尚未准备好使用2.0（尚未发布）

TensorFlow 2.0已发布，但仍不支持具有TPU支持的Colab。一些Google员工说，在2.1发布之后就可以了。

我们在跟踪此进度时遇到问题；https://github.com/huan/tensorflow-handbook-tpu/issues/1#issuecomment-552807573

我的老答案

Googler Wolff证实，我们尚无法在带有TPU的Colab中使用TF 2.0（报告于2019年4月15日）：

您将通过Colab分配的TPU正在运行TF1.x。在Jupyter VM上安装每晚2.0点时，它不会更改TPU。您最终会在Jupyter实例上运行的内容与TPU所拥有的内容之间不匹配。

并且根据https://github.com/tensorflow/tensorflow/issues/24412，对TensorFlow 2.0的TPU支持尚未完成。

解决方案将是监视上述问题并等待TF 2.0发布。