时间戳

首页 > 解决方案 > TPU 卡在 Google Colab 中

python - TPU 卡在 Google Colab 中

问题描述

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

我从 Tensorflow 教程中获得灵感。

它按应有的方式初始化 TPU，它似乎运行完美，直到它到达训练的第一个 epoch，然后它就停止了。

什么都没有发生，它不会中断，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%。

标签： pythontensorflowkerasgoogle-colaboratorytpu

解决方案

抱歉回复晚了，我假设您尝试使用 Tensorflow 2.0，请尝试使用 Tensorflow 2.1 或 2.2。只有从 Tensorflow 2.1 Keras 开始，.compile、.fit、.evaluate 和 .predict API 才可用于 TPU。这是 Tensorflow 2.1发行说明。

对于 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 指南了解更多详情。

推荐阅读