Failing to load saved TFBertModel

[sourabhXIII]

TF version: 2.2.0-rc1
transformers version: 2.7.0

import tensorflow as tf
import transformers
print(tf.__version__)
print(transformers.__version__)
MAX_LEN = 10
model_path = 'saved_model/temp_model'
ids = tf.keras.layers.Input((MAX_LEN,), dtype=tf.int32)
mask = tf.keras.layers.Input((MAX_LEN,), dtype=tf.int32)
token_type_ids = tf.keras.layers.Input((MAX_LEN,), dtype=tf.int32)
base_model = transformers.TFBertModel.from_pretrained("bert-base-cased"
, output_hidden_states=False)
base_output = base_model([ids, mask, token_type_ids])
seq_out, _ = base_output[0], base_output[1]
base_model.trainable = False
model = tf.keras.models.Model(inputs=[ids, mask, token_type_ids], outputs=[seq_out])
model.compile(optimizer='adam', loss='binary_crossentropy', metrics=['accuracy'])
print(model.summary())
model.save(model_path)
model = tf.keras.models.load_model(model_path)

Model load fails with the following error:

Traceback (most recent call last):
File "/Users/sourabhmaity/anaconda3/lib/python3.7/site-packages/tensorflow/python/util/nest.py", line 378, in assert_same_structure
expand_composites)
TypeError: The two structures don't have the same nested structure.

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

Second structure: type=list str=[TensorSpec(shape=(None, 10), dtype=tf.int32, name='inputs/0'), TensorSpec(shape=(None, 10), dtype=tf.int32, name='inputs/1'), TensorSpec(shape=(None, 10), dtype=tf.int32, name='inputs/2')]

More specifically: The two namedtuples don't have the same sequence type. First structure type=dict str={'input_ids': TensorSpec(shape=(None, 5), dtype=tf.int32, name='input_ids')} has type dict, while second structure type=list str=[TensorSpec(shape=(None, 10), dtype=tf.int32, name='inputs/0'), TensorSpec(shape=(None, 10), dtype=tf.int32, name='inputs/1'), TensorSpec(shape=(None, 10), dtype=tf.int32, name='inputs/2')] has type list

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
File "temp.py", line 29, in
model = tf.keras.models.load_model(model_path)
File "/Users/sourabhmaity/anaconda3/lib/python3.7/site-packages/tensorflow/python/keras/saving/save.py", line 190, in load_model
return saved_model_load.load(filepath, compile)
File "/Users/sourabhmaity/anaconda3/lib/python3.7/site-packages/tensorflow/python/keras/saving/saved_model/load.py", line 116, in load
model = tf_load.load_internal(path, loader_cls=KerasObjectLoader)
File "/Users/sourabhmaity/anaconda3/lib/python3.7/site-packages/tensorflow/python/saved_model/load.py", line 604, in load_internal
export_dir)
File "/Users/sourabhmaity/anaconda3/lib/python3.7/site-packages/tensorflow/python/keras/saving/saved_model/load.py", line 188, in init
super(KerasObjectLoader, self).init(*args, **kwargs)
File "/Users/sourabhmaity/anaconda3/lib/python3.7/site-packages/tensorflow/python/saved_model/load.py", line 123, in init
self._load_all()
File "/Users/sourabhmaity/anaconda3/lib/python3.7/site-packages/tensorflow/python/keras/saving/saved_model/load.py", line 215, in _load_all
self._finalize_objects()
File "/Users/sourabhmaity/anaconda3/lib/python3.7/site-packages/tensorflow/python/keras/saving/saved_model/load.py", line 506, in _finalize_objects
_finalize_saved_model_layers(layers_revived_from_saved_model)
File "/Users/sourabhmaity/anaconda3/lib/python3.7/site-packages/tensorflow/python/keras/saving/saved_model/load.py", line 677, in _finalize_saved_model_layers
inputs = infer_inputs_from_restored_call_function(call_fn)
File "/Users/sourabhmaity/anaconda3/lib/python3.7/site-packages/tensorflow/python/keras/saving/saved_model/load.py", line 921, in infer_inputs_from_restored_call_function
spec = nest.map_structure(common_spec, spec, spec2)
File "/Users/sourabhmaity/anaconda3/lib/python3.7/site-packages/tensorflow/python/util/nest.py", line 611, in map_structure
expand_composites=expand_composites)
File "/Users/sourabhmaity/anaconda3/lib/python3.7/site-packages/tensorflow/python/util/nest.py", line 385, in assert_same_structure
% (str(e), str1, str2))
TypeError: The two structures don't have the same nested structure.

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

Second structure: type=list str=[TensorSpec(shape=(None, 10), dtype=tf.int32, name='inputs/0'), TensorSpec(shape=(None, 10), dtype=tf.int32, name='inputs/1'), TensorSpec(shape=(None, 10), dtype=tf.int32, name='inputs/2')]

More specifically: The two namedtuples don't have the same sequence type. First structure type=dict str={'input_ids': TensorSpec(shape=(None, 5), dtype=tf.int32, name='input_ids')} has type dict, while second structure type=list str=[TensorSpec(shape=(None, 10), dtype=tf.int32, name='inputs/0'), TensorSpec(shape=(None, 10), dtype=tf.int32, name='inputs/1'), TensorSpec(shape=(None, 10), dtype=tf.int32, name='inputs/2')] has type list
Entire first structure:
{'input_ids': .}
Entire second structure:
[., ., .]

[frontseat-astronaut]

Facing a similar issue with tf 2.2.0-rc3.

def get_model(lr=0.00001):
    inp_bert = tf.keras.layers.Input(shape=(512), dtype="int32")
    bert = TFBertModel.from_pretrained('bert-base-multilingual-cased')(inp_bert)[0]
    doc_encodings = tf.squeeze(bert[:, 0:1, :], axis=1)
    out = tf.keras.layers.Dense(1, activation="sigmoid")(doc_encodings)
    model = tf.keras.Model(inp_bert, out)
    adam = tf.keras.optimizers.Adam(lr=lr)
    model.compile(optimizer=adam, loss="binary_crossentropy", metrics=["accuracy"])
    return model
model = get_model()
model.save("model_name",save_format='tf')
model = tf.keras.models.load_model('model_name')
model.summary()

Output error is:

/usr/local/lib/python3.6/dist-packages/tensorflow/python/util/nest.py in assert_same_structure(nest1, nest2, check_types, expand_composites)
    383                   "Entire first structure:\n%s\n"
    384                   "Entire second structure:\n%s"
--> 385                   % (str(e), str1, str2))
    386 
    387 

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

First structure: type=TensorSpec str=TensorSpec(shape=(None, 512), dtype=tf.int32, name='inputs')

Second structure: type=dict str={'input_ids': TensorSpec(shape=(None, 5), dtype=tf.int32, name='inputs/input_ids')}

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

[Souls362]

change
base_output = base_model([ids, mask, token_type_ids])
to
base_output = base_model.bert([ids, mask, token_type_ids])
should fix

[bilalghanem]

change
base_output = base_model([ids, mask, token_type_ids])
to
base_output = base_model.bert([ids, mask, token_type_ids])
should fix

Thanks @Souls362 .. solves it.

[camraynor]

This worked for me as well with TFBertModel, however, I run into the same issue with TFXLNetModel. TFXLNetModel doesn't seem to have an equivalent to the .bert property/attribute. Does anyone know how to solve this when using TFXLNetModel?

[LysandreJik]

For TFXLNetModel this would be the .transformer attribute, as you can see here

[camraynor]

@LysandreJik thank you! That works perfectly

[wulikai1993]

@LysandreJik How about TFOpenAIGPTLMHeadModel ? I use .transformer attribute, but the output shape become [None, None, 768], while the original output shape of TFOpenAIGPTLMHeadModel is [None, None, 13088]. How to solve it? Thanks a lot!

[LysandreJik]

Well the transformer attribute is the transformer in itself, which has a hidden size of 768. The LM head model has an additional head which is the embedding matrix of, which has a size of 13088.

[wulikai1993]

Yes, I think so too. So how can I save the whole model?

[Jordy-VL]

@Souls362 you are the greatest! I looked way too long for this.
@huggingface folks, please add an extra detailed example on serialization in TF2.
There is seriously some clear documentation missing there.

[Jordy-VL]

change
base_output = base_model([ids, mask, token_type_ids])
to
base_output = base_model.bert([ids, mask, token_type_ids])
should fix

one tip for TFBertSequenceClassification: base_model.bert([ids, mask, token_type_ids])[1]

[Subfly]

change
base_output = base_model([ids, mask, token_type_ids])
to
base_output = base_model.bert([ids, mask, token_type_ids])
should fix

one tip for TFBertSequenceClassification: base_model.bert([ids, mask, token_type_ids])[1]

What is the difference of 0 and 1 in the brackets?

[Nannigalaxy]

change
base_output = base_model([ids, mask, token_type_ids])
to
base_output = base_model.bert([ids, mask, token_type_ids])
should fix

one tip for TFBertSequenceClassification: base_model.bert([ids, mask, token_type_ids])[1]

What is the difference of 0 and 1 in the brackets?

TFBertModel documentation

model returns sequence output and pooled output (for classification)