问题描述
我的情况是在 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)
解决方法
我发现问题发生是因为我通过写入显式切换到图形模式
from tensorflow.python.framework.ops import disable_eager_execution
disable_eager_execution()
之前
with tpu_strategy.scope():
model.fit(...)
虽然我仍然不明白原因,但删除 disable_eager_execution
解决了问题。