How can we test Transformer Models after converting it to TFLite format

I tried converting the Distilbert Model of Huggingface to TFLite format using this script Reference Script

I was able to convert it. After that, I wanted to test it in Python itself using tf.lite.Interpreter but I was not able to figure out the correct input dimension required for the model. Can anyone help me with this?

Colab Notebook

Pinging @jplu, as he’s working on TF

1 Like


Unfortunately I’m not really familiar with TFLite sorry. As a first guess I would suggest to get your features directly from the Tokenizers and see if it works.

1 Like

Hello @jplu
I tested TFdistilbertModel but after conversion, it’s only expecting tensor of 3*5 dimension. The output of the tokenizer is a single-dimensional array.

I also try to set input like this

input_spec = tf.TensorSpec([1, 768], tf.int32)
model._set_inputs(input_spec, training=False)

But when I check model.inputs it was None and I was getting same error related to dimension.

From what I see in your colab, your TFlite model is waiting only one input (input_ids) and you are giving two (input_ids + attention_mask). Also the method _set_inputs doesn’t work with TF >= 2.2. You have to use:

model._saved_model_inputs_spec = None

When doing TFDistilBertModel.from_pretrained('distilbert-base-uncased') the model is built with a fix input, a batch of three sentences of 3 tokens (then 5 with start and end tokens). So if you want it to fit your needs when creating your savedmodel/TFlite model you have to modify this beforehand as I specified above.


Hi @jplu

Is there a way to allow dynamic input size rather than setting a pre-fixed size? Thanks.

Not for now.

Thanks @jplu
By considering your suggestion by setting input shape, I was able to do inference on TFlite models converted from Huggingface.
Here is my Notebook for someone who is facing such an issue.