Colab TPU 上的 RNN 运行速度与本地 CPU 版本相同

Question

我实现了本地版本的 RNN 和 Colab TPU 版本的 RNN（代码如下）。当我执行 Colab TPU 版本（下面的代码）时，训练速度非常慢，就像我在笔记本电脑 CPU.

上的本地版本运行

Colab TPU 是否支持 RNN 网络？

我是不是漏掉了什么？

import tensorflow as tf
import os
from tensorflow.keras import Sequential
from tensorflow.keras.layers import Dense, SimpleRNN

resolver = tf.distribute.cluster_resolver.TPUClusterResolver(tpu='grpc://' + os.environ['COLAB_TPU_ADDR'])
tf.config.experimental_connect_to_cluster(resolver)
tf.tpu.experimental.initialize_tpu_system(resolver)
print("All devices: ", tf.config.list_logical_devices('TPU'))

strategy = tf.distribute.TPUStrategy(resolver)

with strategy.scope():  
  model = Sequential()
  model.add(SimpleRNN(units=32, input_shape=(1,step), activation="relu"))
  model.add(Dense(16, activation="relu"))
  model.add(Dense(1))
  model.compile(loss='mean_squared_error', optimizer='rmsprop')

model.fit(X,y, epochs=50, batch_size=16, verbose=0)

Answer 1

ctrl-f 在 this page 上用于 RNN。如果你能让 RNN 足够静态，它似乎应该可以工作。

通常，动态操作不适用于 TPU，因为它需要为每个新的训练示例重新编译模型图。

Colab TPU 上的 RNN 运行速度与本地 CPU 版本相同

RNN on Colab TPU runs at the same speed as local CPU version

google-colaboratory

tensorflow2.0

recurrent-neural-network

google-cloud-tpu