You are viewing main version, which requires installation from source. If you'd like
regular pip install, checkout the latest stable version (v1.21.4).
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', 'fill-mask', 'text-generation', 'text2text-generation', 'text-classification', 'token-classification',
'multiple-choice', 'object-detection', 'question-answering', 'image-classification', 'image-segmentation', 'masked-im', 'semantic-
segmentation', 'automatic-speech-recognition', 'audio-classification', 'audio-frame-classification', 'automatic-speech-recognition', 'audio-xvector', 'vision2seq-
lm', 'zero-shot-object-detection', 'text-to-image', 'image-to-image', 'inpainting']. 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.