Convert transformer to SavedModel

Hi! I found out that this is common unresolved problem.

So, I need to convert transformers’ DistilBERT to TensorFlows SavedModel format. I've converted it, but I cant inference it.

Conversion code

import tensorflow as tf
from transformers import TFAutoModel, AutoTokenizer
dir = "distilbert_savedmodel"

model = TFAutoModel.from_pretrained('distilbert-base-uncased')
model.save(dir)

Inference code

tokenizer = AutoTokenizer.from_pretrained('distilbert-base-uncased')
encoded = tokenizer.encode('Hello, world!', add_special_tokens=True, return_tensors="tf")
model = tf.keras.models.load_model(dir)
model(encoded)

Error


ValueError: Could not find matching function to call loaded from the SavedModel. Got:
  Positional arguments (1 total):
    * Tensor("inputs:0", shape=(1, 6), dtype=int32)
  Keyword arguments: {'training': False}

Expected these arguments to match one of the following 4 option(s):

Option 1:
  Positional arguments (1 total):
    * {'input_ids': TensorSpec(shape=(None, 5), dtype=tf.int32, name='input_ids')}
  Keyword arguments: {'training': False}

Option 2:
  Positional arguments (1 total):
    * {'input_ids': TensorSpec(shape=(None, 5), dtype=tf.int32, name='input_ids')}
  Keyword arguments: {'training': True}

Option 3:
  Positional arguments (1 total):
    * {'input_ids': TensorSpec(shape=(None, 5), dtype=tf.int32, name='inputs/input_ids')}
  Keyword arguments: {'training': True}

Option 4:
  Positional arguments (1 total):
    * {'input_ids': TensorSpec(shape=(None, 5), dtype=tf.int32, name='inputs/input_ids')}
  Keyword arguments: {'training': False}

Related issues

huggingface/transformers#4004
huggingface/transformers#2135
huggingface/transformers#2021

Please, help me!

1 Like

In pytorch, you could save the model with something like

torch.save(model.state_dict(),‘/content/drive/My Drive/ftmodelname’ )

Then you could create a model using the pre-trained weights

tuned_model = BertForSequenceClassification.from_pretrained(‘bert-base-uncased’,
num_labels=NCLASSES,
output_attentions=True)

and then overwrite its weights from the saved state_dict with

tuned_model.load_state_dict(torch.load(‘/content/drive/My Drive/ftmodelname’,
map_location=torch.device(“cpu”)),
strict=False)

I expect you could do a similar save of the model state_dict using Tensorflow.

How did you end up solving the problem. I came across the same one.

I have the same problem.

I can train a model, evaluate and predict if it’s within the same script, but if I save the model and then load it I get the same kind of error, is it currently impossible to load a saved model with tensorflow/keras?

hugging face has model.save_pretrained() method