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

时间:2019-04-05 19:04:28

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

由于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分钟)。

1 个答案:

答案 0 :(得分:2)

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