tensorflow: load checkpoint - tensorflow

Question

I've been training a model which looks a bit like:
base_model = tf.keras.applications.ResNet50(weights=weights, include_top=False, input_tensor=input_tensor)
for layer in base_model.layers:
layer.trainable = False
x = tf.keras.layers.GlobalMaxPool2D()(base_model.output)
output = tf.keras.Sequential()
output.add(tf.keras.layers.Dense(2, activation='linear'))
output.add(tf.keras.layers.Dense(2, activation='linear'))
output.add(tf.keras.layers.Dense(2, activation='linear'))
output.add(tf.keras.layers.Dense(2, activation='linear'))
output.add(tf.keras.layers.Dense(2, activation='linear'))
return output(x)
I setup checkpoints saving with code like:
cp_callback = tf.keras.callbacks.ModelCheckpoint(
filepath=checkpoint_path,
verbose=1,
save_weights_only=True,
save_freq=batch_size*5)
Yesterday I started a fit to run for 11 epochs. I'm not sure why, but the machine restarted during the 7th epoch. Naturally I want to resume fitting from the start of epoch 7.
The checkpoint code above created three files:
The contents of checkpoint are:
model_checkpoint_path: "checkpoint"
all_model_checkpoint_paths: "checkpoint"
The other two files are binary. I tried to load the checkpoint weights with both:
model.load_weights('./2022-03-16_21-10/checkpoints/checkpoint.data-00000-of-00001')
model.load_weights('./2022-03-16_21-10/checkpoints/')
Both fail with NotFoundError: Unsuccessful TensorSliceReader constructor: Failed to find any matching files.
How can I restore this checkpoint and as a result resume fitting?
I'm using tensorflow 2.4.

Answer 1

These might help: Training checkpoints and tf.train.Checkpoint. According to the documentation, you should be able to load the model using something like this:
model = tf.keras.Model(...)
checkpoint = tf.train.Checkpoint(model)
# Restore the checkpointed values to the `model` object.
checkpoint.restore(save_path)
I am not sure it will work if the checkpoint contains other variables. You might have to use checkpoint.restore(path).expect_partial().
You can also check the content that has been saved (according to the documentation) by Manually inspecting checkpoints :
reader = tf.train.load_checkpoint('./tf_ckpts/')
shape_from_key = reader.get_variable_to_shape_map()
dtype_from_key = reader.get_variable_to_dtype_map()
sorted(shape_from_key.keys())