How to get/extract CONCRETE FUNCTIONS from my model? (In order to convert my model into TFLite)

Hi, i saved my model in these 3 followings ways:

from transformers import ViTFeatureExtractor, ViTForImageClassification
from transformers import AutoModelForImageClassification
from transformers import TFAutoModelForImageClassification

model_directory = "drive/MyDrive/Tirocinio/1ris_vit-stanford-dog-dataset" 
tokenizer = ViTFeatureExtractor.from_pretrained(model_directory)
model = ViTForImageClassification.from_pretrained(model_directory)

pytorch_path = "drive/MyDrive/Tirocinio/ViT_model_pytorch"
model_pytorch = ViTForImageClassification.from_pretrained(pytorch_path)

h5_path = "drive/MyDrive/Tirocinio/ViT_model_tf"
model_h5 = ViTForImageClassification.from_pretrained(h5_path, from_tf=True)

now what i want to do is get the concrete function in order to convert my model into a TFLite model.

can you please tell me how to do it?

@Rocketknight1

Hi @drgnd4, the easiest way to get a concrete function is to define a TF function that does the computation you want like this:

@tf.function
def forward_pass(inputs):
    outputs = model(inputs).logits

Now you can get a concrete instance of that function by calling get_concrete_function() with a specific input tensor or tensors (whose shape(s) will be used to define the concrete function):

concrete_fn = forward_pass.get_concrete_function(input_tensor)

@Rocketknight1 thank you.

You said that these functions could be used to convert the model in HuggingFace formats to TensorFlowLight but how does it work?

What I am missing is: how can I use these functions to convert my model into a TensorFlowLight model?

Hi @dgrnd4, there are instructions here.

Note that TFLite conversion isn’t something we officially support - we don’t guarantee that our models will convert cleanly! I suspect it should be possible in most cases, but you may run into issues with some ops being unsupported in TFLite.