7

由于 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 分钟/epoch）。

4

1 回答 1

19

Colab 与 TensorFlow 2.2（2020 年 3 月更新）

我解决了这个问题后它就可以工作了，这里还有一个 Colab Notebook

使用 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 模型，然后我们就完成了。

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

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 发布。

于 2019-04-15T09:39:01.347 回答