时间戳

首页 > 解决方案 > 如何使 TensorFlow SimpleRNN 静态以在 Colab TPU 上进行训练？

google-colaboratory - 如何使 TensorFlow SimpleRNN 静态以在 Colab TPU 上进行训练？

问题描述

从 TPU 常见问题解答中阅读以下内容：https ://cloud.google.com/tpu/docs/faq

我可以在 Compute Engine 上训练循环神经网络 (RNN) 吗？

在某些配置中，tf.static_rnn() 和 tf.dynamic_rnn() 与当前的 TPU 执行引擎兼容。更一般地说，TPU 支持 tf.while_loop() 和 TensorArray，它们用于实现 tf.dynamic_rnn()。TPU 不支持诸如 CuDNN 之类的专用工具包，因为它们包含特定于 GPU 的代码。在 TPU 上使用 tf.while_loop() 确实需要指定循环迭代次数的上限，以便 TPU 执行引擎可以静态确定内存使用情况。

如何使我的 SimpleRNN 静态或有效以在 Colab TPU 上运行？

Colab TPU 代码

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)

标签： google-colaboratorytensorflow2.0recurrent-neural-networktpugoogle-cloud-tpu

解决方案

推荐阅读