python - TPU が Google Colab で動かなくなる

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.

動作しているように見えますが、1 時間以上経っても何も起こりませんでした。

Google Cloud Platform によると、TPU の使用率は 0% です (自分の TPU で試したとき)。