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

Posted

技术标签:

【中文标题】Colab TPU 上的 RNN 运行速度与本地 CPU 版本相同【英文标题】:RNN on Colab TPU runs at the same speed as local CPU version 【发布时间】:2020-12-13 12:41:49 【问题描述】:

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

【问题讨论】:

【参考方案1】:

this page 上的 ctrl-f 用于 RNN。如果您可以使 RNN 足够静态,它似乎应该可以工作。

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

【讨论】:

如何使 simpleRNN 静态化?是否会使用“unroll=True”通过此链接实现这一点tensorflow.org/api_docs/python/tf/keras/layers/SimpleRNN

以上是关于Colab TPU 上的 RNN 运行速度与本地 CPU 版本相同的主要内容,如果未能解决你的问题,请参考以下文章

colab使用总结

Google Colab TPU 中未实现文件系统方案“[本地]”

在 Colab TPU 上保存模型时速度极慢

在 Colab TPU 上运行 Pytorch 堆叠模型

使用 TPU 运行时在 Google Colab 上训练 Keras 模型时出错

如何在 Google Colab 中查找 TPU 名称和区域?