Why tflite model output shape is different than the original model converted from T5ForConditionalGeneration?

[generic-matrix]

System Info

• transformers version: 4.24.0
• Platform: Linux-5.10.133+-x86_64-with-Ubuntu-18.04-bionic
• Python version: 3.7.15
• Huggingface_hub version: 0.11.1
• PyTorch version (GPU?): 1.12.1+cu113 (True)
• Tensorflow version (GPU?): 2.9.2 (True)
• Flax version (CPU?/GPU?/TPU?): not installed (NA)
• Jax version: not installed
• JaxLib version: not installed
• Using GPU in script?: no
• Using distributed or parallel set-up in script?: no

Who can help?

@patric @anton-l @sanchit-gandhi @Rocketknight1

Information

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

Tasks

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

Reproduction

T5ForConditionalGeneration Model to translate English to German

from transformers import T5TokenizerFast, T5ForConditionalGeneration

tokenizer = T5TokenizerFast.from_pretrained("t5-small")
model = T5ForConditionalGeneration.from_pretrained("t5-small")

input_ids = tokenizer("translate English to German: the flowers are wonderful.", return_tensors="pt").input_ids
outputs = model.generate(input_ids)
print(tokenizer.decode(outputs[0], skip_special_tokens=True))

Output : Die Blumen sind wunderbar.

Input Shape

input_ids.shape

Output : torch.Size([1, 11])

Output Shape

outputs.shape

Output : torch.Size([1, 7])

Save Pretrained model

!mkdir /content/test
model.save_pretrained('/content/test')

Load TFT5Model model from pretrained

from transformers import TFT5Model
t5model = TFT5Model.from_pretrained('/content/test',from_pt=True)
!mkdir /content/test/t5
t5model.save('/content/test/t5')

Convert TFT5Model to TFlite

import tensorflow as tf

saved_model_dir = '/content/test/t5'
!mkdir  /content/test/tflite
tflite_model_path = '/content/test/tflite/model.tflite'

# Convert the model
converter = tf.lite.TFLiteConverter.from_saved_model(saved_model_dir)

converter.experimental_new_converter = True
converter.experimental_new_quantizer = True
converter.experimental_new_dynamic_range_quantizer = True
converter.allow_custom_ops=True

converter.target_spec.supported_ops = [
  tf.lite.OpsSet.TFLITE_BUILTINS, # enable TensorFlow Lite ops.
  tf.lite.OpsSet.SELECT_TF_OPS # enable TensorFlow ops.
]
converter.optimizations = [tf.lite.Optimize.DEFAULT]
tflite_model = converter.convert()
#print(tflite_model)
print(type(tflite_model))


# Save the model
with open(tflite_model_path, 'wb') as f:
    f.write(tflite_model)

Load The TFLite model

import numpy as np
import tensorflow as tf

tflite_model_path = '/content/test/tflite/model.tflite'
# Load the TFLite model and allocate tensors
interpreter = tf.lite.Interpreter(model_path=tflite_model_path)

interpreter.resize_tensor_input(0, [1,5], strict=True)
interpreter.resize_tensor_input(1, [1,5], strict=True)
interpreter.resize_tensor_input(2, [1,5], strict=True)
interpreter.resize_tensor_input(3, [1,5], strict=True)
interpreter.allocate_tensors()

# Get input and output tensors
input_details = interpreter.get_input_details()
output_details = interpreter.get_output_details()

input_shape = input_details[0]['shape']

#print the output
input_data = np.array(np.random.random_sample((input_shape)), dtype=np.int64)
interpreter.set_tensor(input_details[0]['index'], input_data)
interpreter.invoke()
output_data = interpreter.get_tensor(output_details[0]['index'])

Get The Output Shape

print(output_data.shape)

Expected behavior

print(output_data.shape)

results in

**Output : (1, 8, 5, 64)

Expected something like : (1, 7)**

Can someone let me know where am I going wrong ?

The output shape of the tflite model is completely different from the T5ForConditionalGeneration model

[sgugger]

You reloaded your model in a TFT5Model, which is not the same as T5ForConditionalGeneration: it's the base model without the decoder.

[github-actions[bot]]

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the contributing guidelines are likely to be ignored.

[generic-matrix]

Hello @sgugger ,

Thank you for the update.

Is there any way to convert T5ForConditionalGeneration to TFlite model taking the docs below into consideration ?

https://www.tensorflow.org/api_docs/python/tf/lite/TFLiteConverter

[sanchit-gandhi]

cc @gante @Rocketknight1 the TF experts might be able to help here!

[gante]

Hey @generic-matrix 👋 you probably want to export the entire generation function (which wraps the model), not just the model itself. Look at this test example :)

[nyadla-sys]

@generic-matrix please see the sample notebook converting from TFWhisperForConditionalGeneration to tflite

[github-actions[bot]]

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the contributing guidelines are likely to be ignored.