Question

我正在尝试使用GoogleColab在TPU上运行一些代码。

我从Tensorflow教程中汲取了灵感。

它会按预期的方式初始化TPU，它似乎可以完美运行，直到达到训练的第一个纪元，然后才停止。

什么也不会发生，它不会中断，RAM不会满，但是它永远不会消失。

我已经多次重启环境，但没有任何改变。

autoencoder.fit(
    dataset.batch(1024),
    epochs=100,
    steps_per_epoch=200,
    verbose=1,
    callbacks=[ModelCheckpoint('weights.{epoch:02d}-{loss:.2f}.hdf5', monitor='val_acc', verbose=1, save_best_only=True, mode='max')]
)

模型的定义：

tf.logging.set_verbosity(tf.logging.INFO)
TPU_WORKER = 'grpc://' + os.environ['COLAB_TPU_ADDR']
print(f"TPU: {TPU_WORKER}")
cluster = tf.contrib.cluster_resolver.TPUClusterResolver(TPU_WORKER)
tf.contrib.distribute.initialize_tpu_system(cluster)
strategy = tf.contrib.distribute.TPUStrategy(cluster)
with strategy.scope():
    bvae = ResNetAutoEncoder()
    autoencoder = bvae.ae
    autoencoder.compile(
        optimizer=optimizers.Adam(),
        loss='mean_absolute_error'
    )

我看到以下内容：

Epoch 1/100
W0728 18:42:12.563039 139622569146240 deprecation.py:323] From /usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/engine/training_distributed.py:411: Variable.load (from tensorflow.python.ops.variables) is deprecated and will be removed in a future version.
Instructions for updating:
Prefer Variable.assign which has equivalent behavior in 2.X.

它似乎正在工作，但是即使在一个多小时后，也没有任何反应。

根据Google Cloud Platform（当我尝试使用自己的TPU时），TPU的使用率为0％。

Answer 1

很抱歉，我想您使用Tensorflow 2.0进行了较晚的答复，请尝试使用Tensorflow 2.1或2.2。仅从Tensorflow 2.1 Keras .compile，.fit，.evaluate和.predict API可用于TPU。这是Tensorflow 2.1 release notes。

对于Tensorflow 2.1+，初始化TPUStrategy的代码为：

TPU_WORKER = 'grpc://' + os.environ['COLAB_TPU_ADDR'] # for colab use TPU_NAME if in GCP.

resolver = tf.distribute.cluster_resolver.TPUClusterResolver(TPU_WORKER)
tf.config.experimental_connect_to_cluster(resolver)
tf.tpu.experimental.initialize_tpu_system(resolver)
strategy = tf.distribute.experimental.TPUStrategy(resolver)

阅读TPUStrategy guide了解更多详情。

TPU卡在Google Colab中

1 个答案: