I am fine tuning a classification BERT model using transformers 4.5.1 and tensorflow 2.4.0. I am using SST2 dataset (using TFRecord format).
I am using the simple way:
model = TFBertForSequenceClassification.from_pretrained('bert-base-multilingual-uncased',num_labels=2)
model.summary()
Model: "tf_bert_for_sequence_classification_1"
_________________________________________________________________
Layer (type) Output Shape Param #
=================================================================
bert (TFBertMainLayer) multiple 167356416
_________________________________________________________________
dropout_75 (Dropout) multiple 0
_________________________________________________________________
classifier (Dense) multiple 1538
=================================================================
Total params: 167,357,954
Trainable params: 167,357,954
Non-trainable params: 0
The model is trained in the following way:
tf.keras.backend.clear_session()
learning_rate=5e-5
epsilon=1e-8
# loss
loss = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)
# metric
metric = tf.keras.metrics.SparseCategoricalAccuracy('accuracy')
# optimizer
optimizer = tf.keras.optimizers.Adam(learning_rate=learning_rate, epsilon=epsilon)
model._name = 'tf_bert_classification'
# compile Keras model
model.compile(optimizer=optimizer,
loss=loss,
metrics=[metric])
model.fit(train_data,
epochs=1,
steps_per_epoch=5,
validation_data=eval_data,
validation_steps=1)
The data looks like that: train_data
<BatchDataset shapes: ({input_ids: (None, 128), attention_mask: (None, 128), token_type_ids: (None, 128)}, (None,)), types: ({input_ids: tf.int32, attention_mask: tf.int32, token_type_ids: tf.int32}, tf.int64)>
We can iterate over the data to check again the shape:
for i in eval_data:
print(i[0]['input_ids'].shape)
print(i[0]['attention_mask'].shape)
print(i[0]['token_type_ids'].shape)
print(i[1].shape)
break
(32, 128)
(32, 128)
(32, 128)
(32,)
The model can be trained with a small amount of data. We can also do prediction:
model.predict(eval_data)
FSequenceClassifierOutput(loss=None, logits=array([[-0.9976097 , 0.05450067],
[-0.99764097, 0.05462598],
[-0.9975417 , 0.05509752],
...,
[-0.9975057 , 0.05470059],
[-0.99738246, 0.0544476 ],
[-0.9978782 , 0.05457467]], dtype=float32), hidden_states=None, attentions=None)
I can now save the model and reload it:
model.save('bert_default')
reconstructed_model = tf.keras.models.load_model('bert_default')
reconstructed_model.summary()
Model: "tf_bert_classification"
_________________________________________________________________
Layer (type) Output Shape Param #
=================================================================
bert (TFBertMainLayer) multiple 167356416
_________________________________________________________________
dropout_75 (Dropout) multiple 0
_________________________________________________________________
classifier (Dense) multiple 1538
=================================================================
Total params: 167,357,954
Trainable params: 167,357,954
Non-trainable params: 0
_________________________________________________________________
Now if I try to do prediction with the same data as before I get the following issue:
reconstructed_model.predict(eval_data)
ValueError: Could not find matching function to call loaded from the SavedModel. Got:
Positional arguments (11 total):
* {'input_ids': <tf.Tensor 'input_ids_1:0' shape=(None, 128) dtype=int32>, 'attention_mask': <tf.Tensor 'input_ids:0' shape=(None, 128) dtype=int32>, 'token_type_ids': <tf.Tensor 'input_ids_2:0' shape=(None, 128) dtype=int32>}
* None
* None
* None
* None
* None
* None
* None
* None
* None
* False
Keyword arguments: {}
Expected these arguments to match one of the following 2 option(s):
Option 1:
Positional arguments (11 total):
* {'input_ids': TensorSpec(shape=(None, 5), dtype=tf.int32, name='input_ids/input_ids')}
* None
* None
* None
* None
* None
* None
* None
* None
* None
* True
Keyword arguments: {}
Option 2:
Positional arguments (11 total):
* {'input_ids': TensorSpec(shape=(None, 5), dtype=tf.int32, name='input_ids/input_ids')}
* None
* None
* None
* None
* None
* None
* None
* None
* None
* False
Keyword arguments: {}
I see some Warning when saving the model:
WARNING:tensorflow:The parameter `return_dict` cannot be set in graph mode and will always be set to `True`.
WARNING:absl:Found untraced functions such as embeddings_layer_call_and_return_conditional_losses, embeddings_layer_call_fn, encoder_layer_call_and_return_conditional_losses, encoder_layer_call_fn, pooler_layer_call_and_return_conditional_losses while saving (showing 5 of 1055). These functions will not be directly callable after loading.
WARNING:absl:Found untraced functions such as embeddings_layer_call_and_return_conditional_losses, embeddings_layer_call_fn, encoder_layer_call_and_return_conditional_losses, encoder_layer_call_fn, pooler_layer_call_and_return_conditional_losses while saving (showing 5 of 1055). These functions will not be directly callable after loading.
INFO:tensorflow:Assets written to: ert_default/assets
Is it the reason while the shape is changed from (-1, 128) to (-1, 5) ?
When training the model I see other warning but people as saying they should be ignire:
WARNING:tensorflow:The parameters `output_attentions`, `output_hidden_states` and `use_cache` cannot be updated when calling a model.They have to be set to True/False in the config object (i.e.: `config=XConfig.from_pretrained('name', output_attentions=True)`).
WARNING:tensorflow:The parameter `return_dict` cannot be set in graph mode and will always be set to `True`.
How can I fix this issue that appear only when the model is saved ? Should I set some parameters in the config ?
[Edit]
Same issue with tensorflow 2.5.0-rc3
Using Huggingface save/load is working fine
model.save_pretrained('bert_input')
reconstructed_model = TFBertForSequenceClassification.from_pretrained('bert_input')
But this is not what I need for tf.serving.
It seems some class/function definition are not saved with the model so reloading the model will not work. Maybe there are some paramters to set to have this working ?