Isues with saving and loading tensorflow model which uses hugging face transformer model as its first layer

user_13207502 0
Viewed 1714

Hi I am having some serious problems saving and loading a tensorflow model which is combination of hugging face transformers + some custom layers to do classfication. I am using the latest Huggingface transformers tensorflow keras version. The idea is to extract features using distilbert and then run the features through CNN to do classification and extraction. I have got everything to work as far as getting the correct classifications.

The problem is in saving the model once trained and then loading the model again.

I am using tensorflow keras and tensorflow version 2.2

Following is the code to design the model, train it, evaluate it and then save and load it


    bert_config = DistilBertConfig(dropout=0.2, attention_dropout=0.2, output_hidden_states=False)
    bert_config.output_hidden_states = False
    transformer_model = TFDistilBertModel.from_pretrained(DISTIL_BERT, config=bert_config)

    input_ids_in = tf.keras.layers.Input(shape=(BERT_LENGTH,), name='input_token', dtype='int32')
    input_masks_in = tf.keras.layers.Input(shape=(BERT_LENGTH,), name='masked_token', dtype='int32')

    embedding_layer = transformer_model(input_ids_in, attention_mask=input_masks_in)[0]
    x = tf.keras.layers.Bidirectional(
        tf.keras.layers.LSTM(50, return_sequences=True, dropout=0.1,
                             recurrent_dropout=0, recurrent_activation="sigmoid",
                             unroll=False, use_bias=True, activation="tanh"))(embedding_layer)
    x = tf.keras.layers.GlobalMaxPool1D()(x)

    outputs = []
    # lots of code here to define the dense layers to generate the outputs
    # .....
    # .....

    model = Model(inputs=[input_ids_in, input_masks_in], outputs=outputs)
    for model_layer in model.layers[:3]:
        logger.info(f"Setting layer {model_layer.name} to not trainable")
        model_layer.trainable = False
    rms_optimizer = RMSprop(learning_rate=0.001)
    model.compile(loss=SigmoidFocalCrossEntropy(), optimizer=rms_optimizer)

    # the code to fit the model (which works)
    # then code to evaluate the model (which also works)

    # finally saving the model. This too works.
    tf.keras.models.save_model(model, save_url, overwrite=True, include_optimizer=True, save_format="tf")

However, when I try to load the saved model using the following

    tf.keras.models.load_model(
            path, custom_objects={"Addons>SigmoidFocalCrossEntropy": SigmoidFocalCrossEntropy})

I get the following load error


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

First structure: type=TensorSpec str=TensorSpec(shape=(None, 128), 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, 128), dtype=tf.int32, name='inputs')" is not
Entire first structure:
.
Entire second structure:
{'input_ids': .}

I believe the issue is because TFDistilBertModel layer can be called using a dictionary input from DistilBertTokenizer.encode() and that happens to be the first layer. So the model compiler on load expects that to be the input signature to the call model. However, the inputs defined to the model are two tensors of shape (None, 128)

So how do I tell the load function or the save function to assume the correct signatures?

huggingface-transformers keras tensorflow machine-learning
2 Answers

I solved the issue.

The issue was the object transformer_model in the above code is itself not a layer. So if we want to embed it inside another keras layer we should use the internal keras layer that is wrapped in the model

So changing the line

embedding_layer = transformer_model(input_ids_in, attention_mask=input_masks_in[0]

to

embedding_layer = transformer_model.distilbert([input_ids_in, input_masks_in])[0]

makes everything work. Hope this helps someone else. Took a long time to debug through tf.keras code to figure this one out although in hindsight it is obvious. :)

user_132075020

I suffered the same problem, casually, yesterday. My solution is very similar to yours, I supposed that the problem was due to how tensorflow keras processes custom models so, the idea was to use the layers of the custom model inside my model. This has the advantage of not calling explicitly the layer by its name (in my case, it is useful for easy building more generic models using different pretrained encoders):

sent_encoder = getattr(transformers, self.model_name).from_pretrained(self.shortcut_weights).layers[0]

I don't explored all the models of HuggingFace, but a few that I tested seem to be a custom model with only one custom layer.

Your solution also works like a charm, in fact, both solutions are the same if "distilbert" references to ".layers[0]".

user_127498440
Related