Input shape fixed at 1x5 when converting transformers to tflite

[androbada525]

System Info

Transformers version: 2.8.2 Python version: 3.7.14 on linux Platform: Linux (Google Colab)

Who can help?

@patil-suraj, @patrickvonplaten, @Rocketknight1

Information

• [X] The official example scripts
• [ ] My own modified scripts

Tasks

• [X] An officially supported task in the examples folder (such as GLUE/SQuAD, ...)
• [ ] My own task or dataset (give details below)

Reproduction

Hi,

I am following this tutorial written by the Hugging Face team to convert GPT2 to tflite: https://towardsdatascience.com/on-device-machine-learning-text-generation-on-android-6ad940c00911

As per the tutorial, the generated tflite file should have an input shape of 1x64. However, the input shape turns out as 1x5. There is a Google Colab notebook linked in the tutorial that you can refer to: https://colab.research.google.com/drive/18JPzizAH995pd0iFWx4Xdf-sqjmZwHUD

This is the script that was used in the tutorial for conversion (this script is also in the notebook):

import tensorflow as tf
from transformers import TFGPT2LMHeadModel

model = TFGPT2LMHeadModel.from_pretrained('gpt2')

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

print(model.inputs)
print(model.outputs)

converter = tf.lite.TFLiteConverter.from_keras_model(model)

converter.optimizations = [tf.lite.Optimize.DEFAULT]
converter.target_spec.supported_types = [tf.float16]

tflite_model = converter.convert()

open("gpt2-fp16.tflite", "wb").write(tflite_model)

Notice in the script that the input shape of 1x64 is defined at the beginning

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

However, the input shape of the generated tflite is 1x5.

The input shape can be checked using a website like Netron or by running the following code:

import numpy as np
import tensorflow as tf
from PIL import Image

from os import listdir
from os.path import isfile, join

from random import choice, random

interpreter = tf.lite.Interpreter(model_path="gpt2-fp16.tflite")

interpreter.allocate_tensors()

input_details = interpreter.get_input_details()
output_details = interpreter.get_output_details()

input_shape = input_details[0]['shape']
print(f"Required input shape: {input_shape}")
output_shape = output_details[0]['shape']
print(f"Required output shape: {output_shape}")

It's not just GPT2 that produces an input shape of 1x5. I also tried converting t5-small to tflite and got the same input shape of 1x5. The tflite files for GPT2 on Hugging Face have an input shape of 1x64, though.

The input to GPT-2 can be up to 1024 tokens, and yet the token context length is somehow fixed at 5. A similar issue is present on StackOverflow where the user exported GPT2 as a TF SavedModel and then further to ONNX and TF.js. In both cases, the input shape was 1x5.

I also tried performing the conversion with TFLite Converter v1 API as suggested here:

converter = tf.compat.v1.lite.TFLiteConverter.from_saved_model(
        saved_model_dir, input_arrays=['inputA'], input_shapes={'inputA': [1, 640, 640, 1]})

However, using the v1 API still produced the 1x5 input shape. There was another suggestion, here, to convert the model to SavedModel first and then set the input shape, followed by calling concrete_functions

model = tf.saved_model.load(export_dir)
concrete_func = model.signatures[
  tf.saved_model.DEFAULT_SERVING_SIGNATURE_DEF_KEY]
concrete_func.inputs[0].set_shape([None, 1280, 720, 3])
converter = TFLiteConverter.from_concrete_functions([concrete_func])
...

When using the concreate_functions solution to set the input shape, I got the following error:

InvalidArgumentError: Dimension 1 in both shapes must be equal, but are 5 and 64. Shapes are [?,5] and [1,64].

The error is resolved if I use: concrete_func.inputs[0].set_shape([1,5])

I could not check the input shape accepted by the saved model but from the error above we can get an idea that the SavedModel also uses the 1x5 shape.

I used this code to save the model:

import tensorflow as tf
from transformers import TFGPT2LMHeadModel

model = TFGPT2LMHeadModel.from_pretrained('gpt2')
model.save('gpt2')

Can someone suggest how I can set the input shape to 1x64? Thanks for your time, I appreciate it! :)

Expected behavior

The input shape of the generated tflite file should be 1x64 because that's what we are explicitly defining it as. However, both in the cases of T5 and GPT2, the input shape does not change from 1x5

[LysandreJik]

I wonder if @gante @Rocketknight1 have some insights

[Rocketknight1]

Yes! There are some workarounds you could use @androbada525, such as saving the model as SavedModel and making your TFLite model with the from saved model constructor instead, which will give you the correct signature.

However, from_keras_model should also work, and we have a solution for that exact issue in development right now. I'll ping you in this issue when it's ready!

[Rocketknight1]

Hi @androbada525, this should now be fixed. However, to use the fix you will need to install transformers from main until we release the next version. To do that, replace pip install transformers with pip install --upgrade git+https://github.com/huggingface/transformers.git. After our next release you can change your code back to just pip install transformers.

If this doesn't resolve your problem, feel free to comment and re-open this issue!