Model Export to liteRT #21674

pctablet505 · 2025-09-17T09:21:58Z

This pull request adds support for exporting Keras models to the LiteRT (TFLite) format, along with improvements to input signature handling and export utility documentation. The changes ensure that LiteRT export is only available when TensorFlow is installed, update the export API and documentation, and enhance input signature inference for various model types.

LiteRT Export Support:

Added conditional import of LitertExporter and export_litert in keras/src/export/__init__.py, making LiteRT export available only if TensorFlow is installed.
Updated the Model.export method to support the "litert" format, including new options for LiteRT export and user-facing documentation and example. Raises an informative error if TensorFlow is not installed. [1] [2] [3] [4]
Registered litert as a lazy module in keras/src/utils/module_utils.py for dynamic import support.

Input Signature and Export Utilities:

Improved documentation and logic in get_input_signature to clarify behavior for different model types and ensure correct input signature construction for export. [1] [2]
Enhanced _infer_input_signature_from_model to handle flexible batch dimensions and ensure compatibility with downstream exporters, always returning a flat list of input specs.

Introduces a custom LiteRTExporter for exporting models to TFLite format, bypassing the standard TFLiteConverter. Updates the export API and documentation to support the new 'lite_rt' format, and adds relevant options for custom ops, select TF ops, and optimizations.

Replaces the custom MLIR-based TFLite conversion logic in LiteRTExporter with direct use of the standard TFLiteConverter. Also improves input signature handling for tf.function tracing and updates imports accordingly.

Moved imports of get_input_signature and make_tf_tensor_spec inside functions in saved_model.py to prevent circular imports. Updated EXPORT_FORMATS in export_utils.py to use string references instead of direct imports.

Adds max_sequence_length parameter to input signature generation for sequence models, bounding sequence length for transformer-like architectures. Improves LiteRTExporter with heuristics for complex models, fallback conversion via SavedModel for large models, and exposes max_sequence_length in Model export options. Updates documentation accordingly.

Adds logic to dynamically reduce max_sequence_length for large vocabulary models in export_utils.py to prevent tensor size overflow. In lite_rt_exporter.py, introduces checks and workarounds for models with _DictWrapper issues, and applies memory optimizations for large models during TFLite conversion. These changes improve export reliability and prevent memory errors for models such as Gemma, Llama, and similar architectures.

Removed custom trackable object logic from LiteRTExporter and now save the model directly, simplifying the export process. Also streamlined vocabulary size checks in export_utils to prevent tensor size overflow, removing verbose warnings and redundant comments.

Refactors the TFLite conversion logic in lite_rt_exporter.py to attempt direct conversion first and only fall back to SavedModel if necessary, improving robustness and clarity. Adds a new lite_rt_exporter_simple.py file with a streamlined LiteRTExporter class for direct TFLite export, bypassing complex MLIR conversion paths.

…port

Refactors export_utils and lite_rt_exporter to better detect large vocabulary and Keras-Hub models, applying safer sequence length limits and more robust TFLite conversion paths. Adds heuristics for model type detection, ensures memory safety, and improves handling of TensorFlow introspection issues during export.

Working well with keras

Eliminates the logic for bounding sequence length in model export utilities and related code paths. The max_sequence_length parameter and associated shape bounding for large vocabulary models are removed from export_utils.py and lite_rt_exporter.py. Updates model export documentation accordingly. Adds a comprehensive test script for Keras Hub LiteRT export, verifying numerical accuracy between original and exported models.

fchollet · 2025-10-12T01:52:41Z

keras/src/models/model.py

                    provided, they will be automatically computed.
                - `opset_version`: Optional `int`. Specific to `format="onnx"`.
                    An integer value that specifies the ONNX opset version.
+                - `allow_custom_ops`: Optional `bool`. Specific to


Maybe we should consider putting all the litert args in a single dict to simplify accounting? litert_kwargs.

I think given that every export format already does this through kwargs we should probably be consistent. Either litert_kwargs, onnx_kwargs, tf_save_model_kwargs, etc. Or one final **kwargs that is interpreted per format. No strong preference.

fchollet

Thanks for the PR! The code looks good to me.

Introduces a litert_kwargs parameter for LiteRT model export, allowing users to specify custom export options such as allow_custom_ops, enable_select_tf_ops, and optimizations. This enhances flexibility when exporting models to the LiteRT format.

mattdangerw

Thanks!

keras/src/export/__init__.py

keras/src/export/export_utils.py

keras/src/export/litert.py

mattdangerw · 2025-10-13T20:10:08Z

keras/src/models/model.py

                    provided, they will be automatically computed.
                - `opset_version`: Optional `int`. Specific to `format="onnx"`.
                    An integer value that specifies the ONNX opset version.
+                - `litert_kwargs`: Optional `dict`. Specific to


This seems a bit duplicated. We already have a general passing of kwargs passed to this function to the specific export. Maybe we should just the per format kwargs that are currently supported?

keras/src/export/litert_test.py

keras/src/export/litert.py

mattdangerw · 2025-10-13T20:23:15Z

keras/src/export/litert.py

+                input_shapes = tree.map_structure(
+                    lambda spec: spec.shape, self.input_signature
+                )
+                self.model.build(input_shapes)


I'm not sure we want this? It looks to me like tf saved model export expects the model to be built

keras/keras/src/export/saved_model.py

Lines 151 to 154 in 3137cb0

raise ValueError(

"The layer provided has not yet been built. "

"It must be built before export."

)

and onnx export

keras/keras/src/export/onnx.py

Lines 79 to 82 in 3137cb0

raise ValueError(

"The model provided has never called. "

"It must be called at least once before export."

)

We are just going to make things more confusing if one export format attempt to automatically build but no others do. Let's shoot for consistency.

We are using tf_saved_model as intermediate step, to convert to litert.
We can't expect the model to be fully built/called/traced while calling export.
And making it leave on user, will make the export process complicated from users perspective.

for uniformity, we can change the behaviour in other formats too like onnx.

We can't expect the model to be fully built/called/traced while calling export.

I think we can right? If tf saved model and onnx expect the model to be built, let's just make the same assumption here. It's a good way to start simple. Basically, we'd expect users to hit this error in some cases. Modify their code to call the model on some inputs, and re-export.

Then we could always add this auto build feature as a follow up right? But do it more consistently across export formats.

Moved KerasModelWrapper definition inside LitertExporter for dynamic class creation and removed the old _KerasModelWrapper. Updated import logic for TensorFlow to use module_utils. Improved LiteRT test interpreter selection and simplified test skipping conditions for better backend compatibility.

Updated verbose output in LitertExporter to use io_utils.print_msg instead of print for consistency and better message handling. Warnings about unavailable LiteRT now use the logging module. Improved comments and formatting for clarity.

keras/src/export/export_utils.py

SamanehSaadat · 2025-10-14T21:48:44Z

keras/src/export/litert.py

+    aot_compile_targets=None,
+    **kwargs,
+):
+    """Export the model as a Litert artifact for inference.


I think LiteRT is easier to read than Litert.

keras/src/export/litert.py

SamanehSaadat · 2025-10-14T23:38:10Z

keras/src/export/litert.py

+
+                # Print compilation report if available
+                try:
+                    report = result.compilation_report()


Does this throw an exception if report is not available?

keras/src/export/litert_test.py

Updated all references from LitertExporter to LiteRTExporter in the export module for consistency and clarity. Also corrected related docstrings and messages to use the LiteRT naming.

Improves error messaging in export_utils.py and refines input signature inference logic. Also corrects code block formatting in model.py documentation.

pctablet505 · 2025-09-23T07:31:40Z

keras/src/export/export_utils.py

+# Registry for export formats
+EXPORT_FORMATS = {
+    "tf_saved_model": "keras.src.export.saved_model:export_saved_model",
+    "lite_rt": "keras.src.export.lite_rt_exporter:LiteRTExporter",


Shall it be named 'lite_rt' or 'litert'?

named it as "litert"

pctablet505 · 2025-09-23T07:34:21Z

keras/src/models/model.py

        from keras.src.export import export_saved_model

-        available_formats = ("tf_saved_model", "onnx", "openvino")
+        available_formats = ("tf_saved_model", "onnx", "openvino", "lite_rt")


tflite and lite_rt both may be supported as both generate same format tflite but lite_rt is supposed to be further optimized.

pctablet505 · 2025-10-14T04:53:25Z

keras/src/export/litert.py

+                input_shapes = tree.map_structure(
+                    lambda spec: spec.shape, self.input_signature
+                )
+                self.model.build(input_shapes)


We are using tf_saved_model as intermediate step, to convert to litert.
We can't expect the model to be fully built/called/traced while calling export.
And making it leave on user, will make the export process complicated from users perspective.

for uniformity, we can change the behaviour in other formats too like onnx.

keras/src/export/litert_test.py

keras/src/export/export_utils.py

pctablet505 · 2025-10-21T07:31:44Z

keras/src/export/litert.py

+        self.kwargs = kwargs
+
+    def export(self, filepath):
+        """Exports the Keras model to a TFLite file and optionally performs AOT


LiteRT is just a runtime built on top of TFLite, it generates same old .tflite file.

keras/src/export/export_utils.py

pctablet505 · 2025-10-22T05:28:20Z

keras/src/export/litert.py

+        tflite_model = self._convert_to_tflite(self.input_signature)
+
+        if self.verbose:
+            final_size_mb = len(tflite_model) / (1024 * 1024)


_convert_to_tflite returns the serialized bytes.
len(tflite_model) counts the bytes.

pctablet505 and others added 30 commits August 28, 2025 11:04

Update lite_rt_exporter.py

631850e

Update export_utils.py

f5aa72e

Refactor LiteRTExporter to simplify TFLite conversion

2b952d6

Replaces the custom MLIR-based TFLite conversion logic in LiteRTExporter with direct use of the standard TFLiteConverter. Also improves input signature handling for tf.function tracing and updates imports accordingly.

Refactor import structure to avoid circular dependencies

8f81dd5

Moved imports of get_input_signature and make_tf_tensor_spec inside functions in saved_model.py to prevent circular imports. Updated EXPORT_FORMATS in export_utils.py to use string references instead of direct imports.

trying kerashub

011f1d8

Update lite_rt_exporter.py

d0070c6

Update lite_rt_exporter.py

761793f

Update export_utils.py

c219eb1

Update lite_rt_exporter.py

e26ff6b

Merge branch 'keras-team:master' into export

3aca2f6

Merge branch 'export' of https://github.com/pctablet505/keras into ex…

441a778

…port

Update lite_rt_exporter.py

f4b43b4

Update lite_rt_exporter.py

0fe4bd5

Update lite_rt_exporter.py

8c3faa3

Update lite_rt_exporter.py

88b6a6f

Update lite_rt_exporter.py

da13d04

Update lite_rt_exporter.py

f1f700c

Update lite_rt_exporter.py

5944780

Update lite_rt_exporter.py

4404c39

Working well with keras

Update lite_rt_exporter.py

6a119fb

Merge branch 'keras-team:master' into export

4cec7cd

Merge branch 'keras-team:master' into export

3a7fcc4

Delete test_keras_hub_export.py

e1fca24

pctablet505 added 5 commits October 6, 2025 14:46

Update litert_test.py

d8236fa

Update litert_test.py

83577be

Update litert_test.py

c53b264

Update litert_test.py

487184d

Update requirements-tensorflow-cuda.txt

374d90b

divyashreepathihalli mentioned this pull request Oct 8, 2025

Model Export to liteRT keras-team/keras-hub#2405

Open

pctablet505 requested a review from amitsrivastava78 October 10, 2025 03:49

fchollet reviewed Oct 12, 2025

View reviewed changes

pctablet505 and others added 3 commits October 13, 2025 10:57

Merge branch 'keras-team:master' into export

e843f7e

Update model.py

d01a4cb

mattdangerw reviewed Oct 13, 2025

View reviewed changes

pctablet505 added 3 commits October 14, 2025 10:20

Update export_utils.py

794d85d

SamanehSaadat reviewed Oct 15, 2025

View reviewed changes

pctablet505 and others added 11 commits October 16, 2025 15:58

typo fix

d2b90eb

set verbose to True by default

191f802

removed unnecessary variable

b736ede

Rename LitertExporter to LiteRTExporter

27f1d07

Updated all references from LitertExporter to LiteRTExporter in the export module for consistency and clarity. Also corrected related docstrings and messages to use the LiteRT naming.

Update litert.py

17dccf2

Update export_utils.py

3e16ab3

Fix input signature inference and doc formatting

efbc6d3

Improves error messaging in export_utils.py and refines input signature inference logic. Also corrects code block formatting in model.py documentation.

Update export_utils.py

7825983

Update litert.py

676a53c

Update litert.py

4b6386e

Update litert_test.py

79f05c8

pctablet505 commented Oct 22, 2025

View reviewed changes

Update litert.py

a22eb65

	raise ValueError(
	"The layer provided has not yet been built. "
	"It must be built before export."
	)

	raise ValueError(
	"The model provided has never called. "
	"It must be called at least once before export."
	)

Model Export to liteRT #21674

Are you sure you want to change the base?

Model Export to liteRT #21674

Uh oh!

Conversation

pctablet505 commented Sep 17, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

fchollet left a comment

Choose a reason for hiding this comment

Uh oh!

mattdangerw left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

pctablet505 commented Sep 17, 2025 •

edited

Loading