Search code examples
pythontensorflowgoogle-colaboratorytpu

Use TPU in Google Colab

I am currently training a neural network with the help of a TPU. I changed the runtime type and initialized the TPU. I have the feeling that it is still not faster. I used https://www.tensorflow.org/guide/tpu. Did I something wrong?

# TPU initialization
resolver = tf.distribute.cluster_resolver.TPUClusterResolver(tpu='grpc://' + os.environ['COLAB_TPU_ADDR'])
tf.config.experimental_connect_to_cluster(resolver)
# This is the TPU initialization code that has to be at the beginning.
tf.tpu.experimental.initialize_tpu_system(resolver)
print("All devices: ", tf.config.list_logical_devices('TPU'))

.
.
.
# experimental_steps_per_execution = 50
model.compile(optimizer=Adam(lr=learning_rate), loss='binary_crossentropy', metrics=['accuracy'], experimental_steps_per_execution = 50)

The summary of my model

enter image description here

Is there anything I still have to consider or adjust?

Solution

  • You need to create TPU strategy:

    strategy = tf.distribute.TPUStrategy(resolver).

    And than use this strategy properly:

    with strategy.scope():
  model = create_model()
  model.compile(optimizer='adam',
                loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
                metrics=['sparse_categorical_accuracy'])