Export a model to TFLite with optimum.exporters.tflite

Summary

Exporting a model to TFLite is as simple as

optimum-cli export tflite --model bert-base-uncased --sequence_length 128 bert_tflite/

Check out the help for more options:

optimum-cli export tflite --help

Exporting a model to TFLite using the CLI

To export a 🤗 Transformers model to TFLite, you’ll first need to install some extra dependencies:

pip install optimum[exporters-tf]

The Optimum TFLite export can be used through Optimum command-line. As only static input shapes are supported for now, they need to be specified during the export.

optimum-cli export tflite --help

usage: optimum-cli <command> [<args>] export tflite [-h] -m MODEL [--task TASK] [--atol ATOL] [--pad_token_id PAD_TOKEN_ID] [--cache_dir CACHE_DIR]
                                                    [--trust-remote-code] [--batch_size BATCH_SIZE] [--sequence_length SEQUENCE_LENGTH]
                                                    [--num_choices NUM_CHOICES] [--width WIDTH] [--height HEIGHT] [--num_channels NUM_CHANNELS]
                                                    [--feature_size FEATURE_SIZE] [--nb_max_frames NB_MAX_FRAMES]
                                                    [--audio_sequence_length AUDIO_SEQUENCE_LENGTH]
                                                    output

optional arguments:
  -h, --help            show this help message and exit

Required arguments:
  -m MODEL, --model MODEL
                        Model ID on huggingface.co or path on disk to load model from.
  output                Path indicating the directory where to store generated TFLite model.

Optional arguments:
  --task TASK           The task to export the model for. If not specified, the task will be auto-inferred based on the model. Available tasks depend on
                        the model, but are among: ['default', 'masked-lm', 'causal-lm', 'seq2seq-lm', 'sequence-classification', 'token-classification',
                        'multiple-choice', 'object-detection', 'question-answering', 'image-classification', 'image-segmentation', 'masked-im', 'semantic-
                        segmentation', 'speech2seq-lm', 'audio-classification', 'audio-frame-classification', 'audio-ctc', 'audio-xvector', 'vision2seq-
                        lm', 'stable-diffusion', 'zero-shot-object-detection']. For decoder models, use `xxx-with-past` to export the model using past key
                        values in the decoder.
  --atol ATOL           If specified, the absolute difference tolerance when validating the model. Otherwise, the default atol for the model will be used.
  --pad_token_id PAD_TOKEN_ID
                        This is needed by some models, for some tasks. If not provided, will attempt to use the tokenizer to guess it.
  --cache_dir CACHE_DIR
                        Path indicating where to store cache.
  --trust-remote-code   Allow to use custom code for the modeling hosted in the model repository. This option should only be set for repositories you trust
                        and in which you have read the code, as it will execute on your local machine arbitrary code present in the model repository.

Input shapes:
  --batch_size BATCH_SIZE
                        Batch size that the TFLite exported model will be able to take as input.
  --sequence_length SEQUENCE_LENGTH
                        Sequence length that the TFLite exported model will be able to take as input.
  --num_choices NUM_CHOICES
                        Only for the multiple-choice task. Num choices that the TFLite exported model will be able to take as input.
  --width WIDTH         Vision tasks only. Image width that the TFLite exported model will be able to take as input.
  --height HEIGHT       Vision tasks only. Image height that the TFLite exported model will be able to take as input.
  --num_channels NUM_CHANNELS
                        Vision tasks only. Number of channels used to represent the image that the TFLite exported model will be able to take as input.
                        (GREY = 1, RGB = 3, ARGB = 4)
  --feature_size FEATURE_SIZE
                        Audio tasks only. Feature dimension of the extracted features by the feature extractor that the TFLite exported model will be able
                        to take as input.
  --nb_max_frames NB_MAX_FRAMES
                        Audio tasks only. Maximum number of frames that the TFLite exported model will be able to take as input.
  --audio_sequence_length AUDIO_SEQUENCE_LENGTH
                        Audio tasks only. Audio sequence length that the TFLite exported model will be able to take as input.