Loading tf.keras model, ValueError: The two structures don't have the same nested structure

Question:

I created a tf.keras model that has BERT and I want to train and save it for further use.
Loading this model is a big issue cause I keep getting error: ValueError: The two structures don't have the same nested structure.

I simplified the model a lot, to see where is the problem exactly. The code is pretty simple:

bert = TFBertModel.from_pretrained("bert-base-german-cased")

model_name = "Model"
txt12_input_ids = tf.keras.layers.Input(shape=(max_length,),  name='txt12_input_ids', dtype='int32')
txt12_mask      = tf.keras.layers.Input(shape=(max_length,),  name='txt12_mask', dtype='int32')
txt12_outputs = bert(txt12_input_ids, txt12_mask).pooler_output

model_K = tf.keras.Model(inputs=(txt12_input_ids,  txt12_mask), outputs=txt12_outputs, name=model_name)
model_K.compile(optimizer=Adam(1e-5), loss="binary_crossentropy", metrics="accuracy")


model_K.save(dir_path+'Prob')
model_2 = tf.keras.models.load_model(dir_path+'Prob')

Some notes before you start replying:

  1. I did specified dtype.

  2. No, I don’t want to save just weights.

  3. I tried to use tf.keras.models.save_model(model_K, dir_path+'Prob') instead and it gives the same error.

And the last thing, I work with tf version: 2.6.0. Does anyone knows how to solve it?

Full error message:

ValueError: The two structures don't have the same nested structure.

First structure: type=tuple str=(({'input_ids': TensorSpec(shape=(None, 5), dtype=tf.int32, name='input_ids/input_ids')}, None, None, None, None, None, None, None, None, False), {})

Second structure: type=tuple str=((TensorSpec(shape=(None, 120), dtype=tf.int32, name='input_ids'), TensorSpec(shape=(None, 120), dtype=tf.int32, name='attention_mask'), None, None, None, None, None, None, None, False), {})

More specifically: Substructure "type=dict str={'input_ids': TensorSpec(shape=(None, 5), dtype=tf.int32, name='input_ids/input_ids')}" is a sequence, while substructure "type=TensorSpec str=TensorSpec(shape=(None, 120), dtype=tf.int32, name='input_ids')" is not
Entire first structure:
(({'input_ids': .}, ., ., ., ., ., ., ., ., .), {})
Entire second structure:
((., ., ., ., ., ., ., ., ., .), {})
Asked By: Nadja

||

Answers:

Check out the issue here on GitHub. It should help you to solve the problem with the shape.

Answered By: DataBach

Rebuilding the model isn’t an option for everyone, so I used the custom_objects argument in keras.models.load_model.

model = keras.models.load_model(
    path.join(dir_path, EXPERIMENT_NAME),
    custom_objects={"TFBertModel": TFBertModel}
)

Some additional documentation on the different save formats and loading models with custom objects can be found on there website, here.

Answered By: Roby