Model Export to liteRT #21674
Model Export to liteRT #21674
Conversation
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.
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.
Introduces multiple tests to verify LiteRT export functionality with AOT compilation, including parameter acceptance, multiple targets, quantization optimizations, fallback behavior when LiteRT is unavailable, and the export_with_aot class method. These tests ensure robust handling of AOT compilation scenarios and graceful fallback when infrastructure is missing.
|
@gemini-code-assist review |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces support for exporting Keras models to the LiteRT (TFLite) format, a valuable feature for deploying models on edge devices. The implementation is comprehensive, covering model building, signature inference, conversion with fallbacks, and optional AOT compilation. The documentation has also been updated accordingly.
I've identified a few issues that should be addressed:
- A critical bug in the model wrapping logic that affects multi-input subclassed models.
- A high-severity issue in the input signature inference that can fail for nested inputs.
- A minor code style issue regarding an import statement.
Fixing these issues will improve the robustness and reliability of the new LiteRT export functionality.
There was a problem hiding this comment.
Code Review
This pull request adds support for exporting Keras models to the LiteRT (TFLite) format. The changes include a new LiteRTExporter, updates to the Model.export API and its documentation, and improvements to input signature inference. The implementation is quite comprehensive, with a robust exporter that handles different model types and falls back to a wrapper-based conversion for complex cases. The test suite is also very thorough.
I found one critical issue with the availability check for the LiteRT export, which could lead to unhelpful error messages for users without TensorFlow installed. I also have a minor suggestion to improve the robustness of the model wrapper for subclassed models.
Changed the LiteRT export availability check to require the TensorFlow backend instead of just TensorFlow installation. Updated error message to instruct users to set the backend to TensorFlow.
Replaces passing a list of positional arguments to the model with unpacking using '*args', ensuring correct argument handling when multiple inputs are provided.
Add new test methods to verify SignatureDef functionality in LiteRT export: - test_signature_def_with_named_model: Tests input name preservation - test_signature_def_with_functional_model: Tests functional model export - test_signature_def_with_multi_input_model: Tests multi-input model handling - test_signature_def_with_multi_output_model: Tests multi-output model handling These tests ensure that exported TFLite models maintain proper input/output names in their SignatureDef and work correctly with signature runners for inference. All tests pass and verify compatibility across different model types (functional, multi-input, multi-output).
|
@gemini-code-assist review |
There was a problem hiding this comment.
Code Review
This pull request introduces support for exporting Keras models to the LiteRT (TFLite) format. The changes are comprehensive, including the core exporter logic, updates to the Model.export API, documentation, and an extensive test suite. The implementation is robust, with fallbacks for different model types and good error handling.
My review focuses on a few key areas for improvement:
- Ensuring that the new LiteRT dependency, which is TensorFlow-specific, does not break Keras for users of other backends (JAX, PyTorch).
- Improving logging consistency within the new exporter module.
- Reducing reliance on internal APIs of dependencies to improve long-term maintainability.
Overall, this is a great addition to Keras's export capabilities.
Replaces direct printing of traceback with formatted traceback output using io_utils.print_msg for better consistency in verbose error reporting.
keras/src/export/litert.py
Outdated
| input_shapes = tree.map_structure( | ||
| lambda spec: spec.shape, self.input_signature | ||
| ) | ||
| self.model.build(input_shapes) |
There was a problem hiding this comment.
Let's not auto-build. Exporting a freshly built model (un-trained, not loaded) with random weights makes no sense. Please remove lines 167-198 and instead throw an error.
keras/src/export/litert.py
Outdated
| elif self.verbose: | ||
| io_utils.print_msg( | ||
| f"Warning: Unknown converter setting '{key}' - ignoring" | ||
| ) |
There was a problem hiding this comment.
else raise a ValueError, we should not ignore anything that is incorrect.
Introduces methods to detect and handle models with dictionary inputs, including input signature inference and adapter model creation for TFLite export. Updates export logic to convert dict input signatures to list format and ensures compatibility with TFLite conversion. Also improves handling of multi-input models and input signature processing.
Introduces multiple tests to verify LiteRT export functionality for models with dictionary inputs, including adapter creation, input signature inference, custom signatures, numerical accuracy, variable sharing, and multi-output support. These tests ensure correct export and inference behavior for dict input models.
Streamlined the ImportError message for LiteRT export to remove backend switching instructions. Also removed the unused 'verbose' parameter from the export_litert call.
Eliminated the _ensure_model_built method and its invocation from LiteRTExporter. The exporter now skips explicit model building before input signature resolution, simplifying the export process.
Eliminates Ahead-of-Time (AOT) compilation functionality from the LiteRT exporter and related tests. The exporter no longer accepts or processes aot_compile_targets, and all AOT-related methods and test cases have been removed for simplification and maintenance. AOT compilation was opensouce earlier, now removed by google as opensource feature.
Moved import statements for export_utils functions to the top of the file and removed redundant inline imports. This improves code clarity and reduces repeated imports within methods.
Replaces assertion with ValueError when the export filepath does not end with '.tflite' in LiteRTExporter. Updates corresponding test to expect ValueError instead of AssertionError for incorrect file extension.
|
@gemini-code-assist review |
There was a problem hiding this comment.
Code Review
This pull request introduces support for exporting Keras models to the LiteRT (TFLite) format. The implementation is comprehensive, handling various model types, including those with dictionary inputs, and includes robust fallback mechanisms for the conversion process. The accompanying tests are thorough and cover a wide range of scenarios.
My review focuses on improving maintainability, debuggability, and adherence to the project's documentation standards. I've identified some unused code, suggested adding more detailed logging to exception handling to aid in debugging conversion failures, and recommended adding comments to a particularly complex method to improve its clarity. I also pointed out a missing code example in a docstring, to align with the Keras API design guidelines.
| def _has_dict_inputs(self): | ||
| """Check if the model expects dictionary inputs. | ||
| Returns: | ||
| bool: True if model has dict inputs, False otherwise. | ||
| """ | ||
| # Check if model.inputs is a dict (most reliable for built models) | ||
| if hasattr(self.model, "inputs") and isinstance( | ||
| self.model.inputs, dict | ||
| ): | ||
| return True | ||
|
|
||
| # Check if _inputs_struct is a dict (for Functional models) | ||
| if hasattr(self.model, "_inputs_struct") and isinstance( | ||
| self.model._inputs_struct, dict | ||
| ): | ||
| return True | ||
|
|
||
| # Check if provided input_signature is a dict | ||
| if self.input_signature is not None: | ||
| if isinstance(self.input_signature, dict): | ||
| return True | ||
| # Check for wrapped dict (Functional model pattern) | ||
| if ( | ||
| isinstance(self.input_signature, (list, tuple)) | ||
| and len(self.input_signature) == 1 | ||
| and isinstance(self.input_signature[0], dict) | ||
| ): | ||
| return True | ||
|
|
||
| return False |
There was a problem hiding this comment.
The method _has_dict_inputs is defined but appears to be unused within the LiteRTExporter class. To improve code clarity and maintainability, this dead code should be removed.
| """Get input signature for model export. | ||
| Args: | ||
| model: A Keras Model instance. | ||
| Returns: | ||
| Input signature suitable for model export (always a tuple or list). | ||
| """ |
There was a problem hiding this comment.
The new docstring for get_input_signature is helpful, but it's missing a code example. According to the Keras API design guidelines, all docstrings should include code examples to demonstrate usage.1
Style Guide References
Footnotes
-
The style guide states: 'All docstrings should include code examples.' ↩
|
|
||
| except Exception: |
There was a problem hiding this comment.
The except Exception: block is very broad and silently falls back to the wrapper-based conversion. For better debuggability, it would be helpful to log the exception that caused the direct conversion to fail. This would provide valuable insight into why the fallback was necessary.
except Exception as e:
io_utils.print_msg(
f"Direct conversion to TFLite failed with error: {e}. "
"Falling back to wrapper-based conversion."
)
return self._convert_with_wrapper(input_signature)| @tf.function | ||
| def __call__(self, *args, **kwargs): | ||
| """The single entry point for the exported model.""" | ||
| # Handle both single and multi-input cases | ||
| if args and not kwargs: | ||
| # Called with positional arguments | ||
| if len(args) == 1: | ||
| return self._model(args[0]) | ||
| else: | ||
| # Multi-input case: Functional models expect a list, | ||
| # not unpacked positional args | ||
| if ( | ||
| hasattr(self._model, "inputs") | ||
| and len(self._model.inputs) > 1 | ||
| ): | ||
| return self._model(list(args)) | ||
| else: | ||
| return self._model(*args) | ||
| elif kwargs and not args: | ||
| # Called with keyword arguments | ||
| if len(kwargs) == 1 and "inputs" in kwargs: | ||
| # Single input case | ||
| return self._model(kwargs["inputs"]) | ||
| else: | ||
| # Multi-input case - convert to list/dict format | ||
| # expected by model | ||
| if ( | ||
| hasattr(self._model, "inputs") | ||
| and len(self._model.inputs) > 1 | ||
| ): | ||
| # Multi-input functional model | ||
| input_list = [] | ||
| missing_inputs = [] | ||
| for input_layer in self._model.inputs: | ||
| input_name = input_layer.name | ||
| if input_name in kwargs: | ||
| input_list.append(kwargs[input_name]) | ||
| else: | ||
| missing_inputs.append(input_name) | ||
|
|
||
| if missing_inputs: | ||
| available = list(kwargs.keys()) | ||
| raise ValueError( | ||
| f"Missing required inputs for multi-input " | ||
| f"model: {missing_inputs}. " | ||
| f"Available kwargs: {available}. " | ||
| f"Please provide all inputs by name." | ||
| ) | ||
|
|
||
| return self._model(input_list) | ||
| else: | ||
| # Single input model called with named arguments | ||
| return self._model(list(kwargs.values())[0]) | ||
| else: | ||
| # Fallback to original call |
There was a problem hiding this comment.
The __call__ method in KerasModelWrapper is quite complex, as it handles many different ways a Keras model can be called (e.g., *args, **kwargs, single/multi-input). While this is necessary for broad compatibility, the logic is difficult to follow. Adding comments to explain the different cases, especially the logic for multi-input functional models, would significantly improve the readability and future maintainability of this code.
|
|
||
| except Exception: |
There was a problem hiding this comment.
The except Exception: block here is very broad and silently tries the next strategy. This can make debugging difficult. It would be more informative to log which strategy failed and why. Consider collecting the error messages from each failed strategy and including them in the final RuntimeError if all strategies fail.
8 participants
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
exportAPI and documentation, and enhance input signature inference for various model types.LiteRT export Design Doc
LiteRT Export Support:
LitertExporterandexport_litertinkeras/src/export/__init__.py, making LiteRT export available only if TensorFlow is installed.Model.exportmethod 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]litertas a lazy module inkeras/src/utils/module_utils.pyfor dynamic import support.Input Signature and Export Utilities:
get_input_signatureto clarify behavior for different model types and ensure correct input signature construction for export. [1] [2]_infer_input_signature_from_modelto handle flexible batch dimensions and ensure compatibility with downstream exporters, always returning a flat list of input specs.