如何在 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 分钟/epoch与 GPU 上的 2 分钟/纪元相比）。

Answer 1

使用 TensorFlow 2.2 的 Colab（2020 年 3 月更新）

在我修复了这个issue 后它可以工作了，here 上还有一个 Colab 笔记本

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

使用 TensorFlow 2.0 将 Keras 模型与 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])

我们创建一个分布式解析器，然后为解析器创建策略，然后使用strategy.scope()，我们创建我们的 Keras 模型，然后我们就完成了。

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

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

但是，请注意 Colab 存在一些需要修复的环境问题，因此它可能还无法在 Colab 中运行。

Colab TPU 尚未准备好 2.0（还）

TensorFlow 2.0 已经发布，但支持 TPU 的 Colab 尚不支持。有谷歌人说2.1发布后就可以了。

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

我的旧答案

Googler Wolff 证实，我们还不能在 Colab 中使用 TPU 中的 TF 2.0（报告于 2019 年 4 月 15 日）：

您将通过 Colab 分配的 TPU 正在运行 TF 1.x。在 Jupyter VM 上安装 nightly 2.0 pip 时，它不会更改 TPU。您最终会发现 Jupyter 实例上运行的内容与 TPU 所拥有的内容不匹配。

根据https://github.com/tensorflow/tensorflow/issues/24412，TPU 对 TensorFlow 2.0 的支持还不完整。

解决方法是监控上述问题，等待 TF 2.0 发布。