在TPU上进行训练时如何在tf keras中节省模型权重？

Question

我想保存我在kaggle的TPU上训练过的模型的权重，但出现错误。它是一个已知的错误，一种解决方法是将权重从TPU节省下来。此错误还有其他解决方法吗？

Answer 1

从Tensorflow 2.3和tf每晚开始，您可以将tf.saved_model.LoadOptions与tf.saved_model.LoadOptions =“ / job：localhost”一起使用，以在本地文件系统上保存/恢复TPU模型：

model_dir = "./mnist_model"

localhost_save_option = tf.saved_model.SaveOptions(experimental_io_device="/job:localhost")
model.save(model_dir, options=localhost_save_option)

# Restore the weights
model2 = tf.keras.models.load_model(model_dir, options=localhost_save_option)