在 Colab TPU 上保存模型时速度极慢

Question

我的情况是在Colab TPU环境下保存模型特别慢

我在使用 checkpoint 回调时第一次遇到这个问题，导致训练卡在第一个 epoch 的末尾。

然后，我尝试取消回调并使用 model.save_weights() 保存模型，但没有任何改变。通过使用 Colab 终端，我发现保存速度约为 100k，持续 5 分钟。

Tensorflow 版本 = 2.3

我的模型拟合代码在这里：

with tpu_strategy.scope(): # creating the model in the TPUStrategy scope means we will train the model on the TPU

    Baseline = create_model()
    checkpoint = keras.callbacks.ModelCheckpoint('baseline_{epoch:03d}.h5', 
                                 save_weights_only=True, save_freq="epoch")


    hist = model.fit(get_train_ds().repeat(), 
                steps_per_epoch = 100,
                epochs = 5,
                verbose = 1,
                callbacks = [checkpoint])

    model.save_weights("epoch-test.h5", overwrite=True)

Answer 1

我发现问题的发生是因为我通过写

明确切换到图形模式

from tensorflow.python.framework.ops import disable_eager_execution
disable_eager_execution()

之前

with tpu_strategy.scope():
    model.fit(...)

虽然我仍然不明白原因，但删除 disable_eager_execution 解决了问题。