残差神经网络模型在 google colab tpu 硬件上运行很慢？答案

【问题标题】：Residual neural network model runs very slowly on google colab tpu hardware?残差神经网络模型在 google colab tpu 硬件上运行很慢？
【发布时间】：2019-06-14 16:39:22
【问题描述】：

我在 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

预计每个 epoch 至少在一分钟内完成（通常最多大约 10 秒），但似乎每个 mini-batch 需要一整分钟才能完成（并且每个 epoch 有很多 mini-batch）。

【问题讨论】：

我找到了一种方法可以让一些 keras resnet 演示代码在 tpu 上快速运行，所以我会从那里向后工作并在解决问题时回答我自己的问题。

标签： keras deep-learning google-colaboratory google-cloud-tpu tpu

【解决方案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。

【讨论】：

【解决方案2】：

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

resolver = tf.distribute.cluster_resolver.TPUClusterResolver()
tf.config.experimental_connect_to_cluster(resolver)
tf.tpu.experimental.initialize_tpu_system(resolver)
strategy = tf.distribute.experimental.TPUStrategy(resolver)

with strategy.scope():
  model = create_model()
  model.compile(optimizer='adam',
                loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
                metrics=['sparse_categorical_accuracy'])

更多信息请关注tpu guide。

【讨论】：