在Google Colab中使用TPU

Question

我目前正在TPU的帮助下训练一个神经网络。我更改了运行时类型并初始化了TPU。我有一种感觉，它还是不快。我使用了https://www.tensorflow.org/guide/tpu。我做错什么了吗？

# 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)

我的模型总结

​

​

还有什么需要我考虑或调整的吗？

Answer 1

您需要创建TPU策略：

strategy = tf.distribute.TPUStrategy(resolver).

然后正确地使用这个策略：

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