TPU застревает в Google Colab

Я пытаюсь запустить код на TPU с помощью GoogleColab.

Я вдохновился учебником Tensorflow.

Он инициализирует TPU, как должен, он работает отлично, пока не достигнет первой эпохи обучения, а затем просто останавливается.

Больше ничего не происходит, не ломается, оперативная память не заполнена, но дальше не идет.

Я уже несколько раз перезапускал среду, но ничего не изменилось.

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.

Вроде работает, но даже через час больше ничего не произошло.

TPU использует 0% в соответствии с Google Cloud Platform (когда я пробовал на своем собственном TPU).


person Julien Séveno-Piltant    schedule 28.07.2019    source источник


Ответы (1)


Извините за поздний ответ, я предполагал, что вы пробовали использовать Tensorflow 2.0, попробуйте Tensorflow 2.1 или 2.2. Только из Tensorflow 2.1 API Keras .compile, .fit, .evaluate и .predict доступны для 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.

person Gagik    schedule 12.06.2020