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

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

I implemented a local version of an RNN and a Colab TPU version of an RNN(code-below). When I execute the Colab TPU version(code-below), the training speed is very slow like my local version running on my laptop's CPU.

Does Colab TPU support RNN networks?

Am I missing something here?

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)

Solution

ctrl-f on this page for RNN. It seems like it should work if you can make the RNN static enough.

In general, dynamic operations don't work well with TPUs since it needs to recompile the model graph for each new training example.