Tensorflow TimeDistributed Wrapped Model Load/Save

Following the tutorial on Tensorflow here, once the following model has been trained:

tensorflow version: 2.17.0

Input Shape: (None, 16, 224, 224, 3)

Which is (BATCH_SIZE, SEQUENCE, 224, 224,3)

net = tf.keras.applications.EfficientNetB0(include_top = False)
net.trainable = False

model = tf.keras.Sequential([
    tf.keras.layers.Rescaling(scale=255),
    tf.keras.layers.TimeDistributed(net),
    tf.keras.layers.Dense(10),
    tf.keras.layers.GlobalAveragePooling3D()
])

What i need assistance with is, how exactly would i save a model built with this architecture?

I tried eg. model.save('model.keras'), which saves the model but the moment i try: loaded_model = tf.keras.models.load_model('model.keras')

i get the following error:

ValueError: Exception encountered when calling TimeDistributed.call().

Cannot convert '16' to a shape.

Arguments received by TimeDistributed.call():
  • args=('<KerasTensor shape=(None, 16, 224, 224, 3), dtype=float32, sparse=False, name=keras_tensor_2462>',)
  • kwargs={'mask': 'None'}

Also Tried different Ways of Saving/Loading the model:

# serialize model to JSON
model_json = model.to_json()
with open("model.json", "w") as json_file:
    json_file.write(model_json)
# serialize weights to HDF5
model.save_weights("model.weights.h5")
print("Saved model to disk")

# load json and create model
json_file = open('model.json', 'r')
loaded_model_json = json_file.read()
json_file.close()
loaded_model = tf.keras.models.model_from_json(loaded_model_json)
# load weights into new model
loaded_model.load_weights("model.weights.h5")
print("Loaded model from disk")

Update:

Managed to get it working by using the following:

model.export('model')

and then importing it using:

loaded_model = tf.saved_model.load('model')

and then using

y = loaded_model .serve(np.expand_dims(sample_video, axis = 0))