Colab TPU上的RNN运行速度与本地CPU版本相同 [英] RNN on Colab TPU runs at the same speed as local CPU version

问题描述

我实现了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)

推荐答案

ctrl-f onthis page用于rnn。如果您可以使RNN足够静态，它似乎应该可以工作。

通常，动态操作不能很好地与TPU配合使用，因为它需要为每个新的训练示例重新编译模型图。

这篇关于Colab TPU上的RNN运行速度与本地CPU版本相同的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！