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

Posted 2023-03-16

【中文标题】在 Colab TPU 上保存模型时速度极慢

【英文标题】：Extremely slow when saving model on Colab TPU

【发布时间】：2021-07-28 20:57:12

【问题描述】：

我的情况是 Colab TPU 环境下保存模型非常慢。

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

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

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)

【问题讨论】：

5 分钟相当长。模型有多大？同样作为健全性检查，model.save_weights("epoch-test.h5", ...) 暗示它将保存到 Colab 机器的文件系统 - 如果将其保存在 GCS 上会更快吗？

@AllenWang 感谢您的回复。该模型约为 190 MB。其实我不知道如何在 GCS 上保存模型，总是尝试在 Colab 机器上保存。所以，我提到的时间是在 Colab 上，而不是 GCS

【参考方案1】：

我发现问题的发生是因为我通过写入显式切换到图形模式

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

之前

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

虽然我还是不明白原因，删除disable_eager_execution解决了这个问题。