残留神经网络模型在Google Colab TPU硬件上运行得非常慢？

Question

我已经在keras的Google Colab上为cifar10数据集建立了一个残差神经网络模型，但是它在TPU硬件上的运行速度非常慢。

我有另一个常规的卷积神经网络，可以在google colab上正常运行。该模型使用keras Sequential API，而残余神经网络使用Functional API，不确定是否是问题所在。我已经尝试更改批次大小，但没有帮助。我程序的链接在下面。

https://colab.research.google.com/github/valentinocc/Keras_cifar10/blob/master/keras_rnn_cifar10.ipynb#scrollTo=7Jc51Dbac2MC

预计每个时期至少要在一分钟内完成（通常最多10秒左右），但似乎每个迷你批次都需要一整分钟才能完成（每个时期有很多迷你批次）。

Answer 1

问题似乎与使用不同的优化器有关。使用tensorflow.keras.optimizers.Adam允许tpu正常运行，而使用tensorflow.train.AdamOptimizer的运行非常慢。

但是，当使用函数fit_generator和ImageDataGenerator对象进行数据扩充时，该问题变得更加难以解决（与仅“ fit”函数相对）。 ImageDataGenerator，keras功能API和TPU硬件似乎不能很好地协同工作。 tf.keras.optimizers.Adam将出现运行时错误，并且tf.train.AdamOptimizer的运行速度与CPU一样快。我认为这里的解决方案是使用带有GPU的另一个框架，或者尝试没有keras的tensorflow。

Answer 2

它似乎表明您的模型不是在TPU硬件上运行，而是在CPU上运行。 为了在Tensorflow Keras模型的TPU上运行训练/预测，您将创建一个TPUStrategy并在该策略范围内编译模型：

nil

有关更多信息，请遵循tpu guide。