Как сохранить веса моделей в tf keras при тренировках на TPU?

Question

Я хочу сохранить веса модели, которую я тренировал на TPU от kaggle, но получаю сообщение об ошибке. Это известная ошибка, и можно обойтись без сохранения весов в ЦП из TPU. Есть ли другой способ обхода этой ошибки?

Answer 1

Начиная с Tensorflow 2.3 и tf-nightly вы можете использовать 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)