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

我实现了RNN的本地版本和RNN的Colab TPU版本（代码如下）。当我执行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)

为RNN启用ctrl-f。如果您能使RNN足够静态，它似乎应该可以工作

---

一般来说，动态操作不适合TPU，因为它需要为每个新的培训示例重新编译模型图。

如何使SimpleRN成为静态的？使用“unroll=True”会从这个链接实现吗