Model Export to liteRT #2405
Model Export to liteRT #2405
Conversation
This reverts commit 62d2484.
This reverts commit de830b1.
Refactored exporter and registry logic for better type safety and error handling. Improved input signature methods in config classes by extracting sequence length logic. Enhanced LiteRT exporter with clearer verbose handling and stricter error reporting. Registry now conditionally registers LiteRT exporter and extends export method only if dependencies are available.
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Summary of Changes
Hello @pctablet505, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!
This pull request introduces a comprehensive and extensible framework for exporting Keras-Hub models to various formats, with an initial focus on LiteRT. The system is designed to seamlessly integrate with Keras-Hub's model architecture, particularly by addressing the unique challenge of handling dictionary-based model inputs during the export process. This enhancement significantly improves the deployability of Keras-Hub models by providing a standardized and robust export pipeline, alongside crucial compatibility fixes for TensorFlow's SavedModel/TFLite export mechanisms.
Highlights
- New Model Export Framework: Introduced a new, extensible framework for exporting Keras-Hub models, designed to support various formats and model types.
- LiteRT Export Support: Added specific support for exporting Keras-Hub models to the LiteRT format, verified for models like gemma3, llama3.2, and gpt2.
- Registry-Based Configuration: Implemented an
ExporterRegistryto manage and retrieve appropriate exporter configurations and exporters based on model type and target format. - Input Handling for Keras-Hub Models: Developed a
KerasHubModelWrapperto seamlessly convert Keras-Hub's dictionary-based inputs to the list-based inputs expected by the underlying Keras LiteRT exporter. - TensorFlow Export Compatibility: Added compatibility shims (
_get_save_specand_trackable_children) to Keras-HubBackbonemodels to ensure proper functioning with TensorFlow's SavedModel and TFLite export utilities. - Automated Export Method Extension: The
Taskclass in Keras-Hub models is now automatically extended with anexportmethod, simplifying the model export process for users.
Using Gemini Code Assist
The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.
Invoking Gemini
You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.
| Feature | Command | Description |
|---|---|---|
| Code Review | /gemini review |
Performs a code review for the current pull request in its current state. |
| Pull Request Summary | /gemini summary |
Provides a summary of the current pull request in its current state. |
| Comment | @gemini-code-assist | Responds in comments when explicitly tagged, both in pull request comments and review comments. |
| Help | /gemini help |
Displays a list of available commands. |
Customization
To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.
Limitations & Feedback
Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.
You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.
Footnotes
-
Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩
There was a problem hiding this comment.
Code Review
This pull request introduces a significant new feature: model exporting to liteRT. The implementation is well-structured, using a modular and extensible registry pattern. However, there are several areas that require attention. The most critical issue is the complete absence of tests for the new export functionality, which is a direct violation of the repository's style guide stating that testing is non-negotiable. Additionally, I've identified a critical bug in the error handling logic within the lite_rt.py exporter that includes unreachable code. There are also several violations of the style guide regarding the use of type hints in function signatures across all new files. I've provided specific comments and suggestions to address these points, which should help improve the robustness, maintainability, and compliance of this new feature.
Introduces the keras_hub.api.export submodule and updates the main API to expose it. The new export module imports various exporter configs and functions from the internal export package, making them available through the public API.
Added ImageClassifierExporterConfig, ImageSegmenterExporterConfig, and ObjectDetectorExporterConfig to the export API. Improved input shape inference and dummy input generation for image-related exporter configs. Refactored LiteRTExporter to better handle model type checks and input signature logic, with improved error handling for input mapping.
Moved the 'import keras' statement to the top of the module and removed redundant local imports within class methods. This improves code clarity and avoids repeated imports.
Deleted the debug_object_detection.py script, which was used for testing object detection model outputs and export issues. This cleanup removes unused debugging code from the repository.
Renames all references of 'LiteRT' to 'Litert' across the codebase, including file names, class names, and function names. Updates exporter registry and API imports to use the new 'litert' naming. Also improves image model exporter configs to dynamically determine input dtype from the model, enhancing flexibility for different input types. Adds support for ImageSegmenter model type detection in the exporter registry.
Refactored InputSpec definitions in exporter configs for improved readability by placing each argument on a separate line. Updated import path in litert.py to import from keras.src.export.litert instead of keras.src.export.litert_exporter.
|
@pctablet505 can you update the Colab to use the changes from this PR? and keep the demo short where you load a model and export it, and reload and verify numerics? |
|
Also if you have addressed the Gemini comments, please respond on it. Is this PR dependent on the PR in Keras repo? - keras-team/keras#21674 |
Simplifies and unifies input signature and dummy input generation for text and image models by removing redundant helper methods and centralizing logic. Updates model building in KerasHubExporter to use input signatures and improves error handling. Refactors LiteRT exporter to use the new parameterized input signature and model building approach, reducing code duplication and improving maintainability.
Removed redundant registry initialization in export_model and clarified model building comments in KerasHubExporter. Switched to using warnings.warn for import errors in models/__init__.py instead of print statements for better error reporting.
|
@gemini-code-assist please review it again |
There was a problem hiding this comment.
Code Review
This is a significant and valuable addition to KerasHub, enabling model exportation to LiteRT. The overall structure is well-designed, following a modular and extensible pattern similar to Hugging Face's Optimum.
I've reviewed this pull request, focusing on correctness, maintainability, and adherence to the repository's style guide. The style guide was instrumental in this review, particularly the sections on Testing and Validation (lines 401-496) and Docstrings and Type Hinting (lines 364-398).
My main findings are:
-
Critical Issue: Missing Tests
- The pull request introduces substantial new functionality for model exporting across several new files (
base.py,configs.py,litert.py,registry.py), but it lacks corresponding tests. - The repository style guide is explicit that "Testing is a non-negotiable part of every contribution" (line 403) and "Every .py file containing logic...must have a corresponding
_test.pyfile" (line 406). - Please add comprehensive unit tests for the new export logic, covering different model types, configurations, and edge cases. This is crucial to ensure the robustness and correctness of this feature.
- The pull request introduces substantial new functionality for model exporting across several new files (
-
Other Findings
- I've also left several inline comments regarding a bug in model type detection, incorrect dtype handling, and violations of the docstring style guide. Please address these to improve code quality and consistency.
Refined dtype extraction logic in image and object model exporter configs to better handle different dtype representations. Updated LiteRT exporter to use Keras io_utils for progress messages and improved verbose flag handling. Added ObjectDetector and ImageSegmenter to export registry model type checks. Enhanced docstrings for clarity and consistency in base exporter classes.
Updates LiteRT exporter and related configs to support dynamic input shapes by default for text and image models, allowing runtime resizing via TFLite's interpreter.resize_tensor_input(). Removes static sequence length inference, adapts input signature logic, and updates tests to verify dynamic shape support and runtime resizing. Also improves multimodal model handling and input mapping for TFLite export.
Adds tests to verify that SignatureDef preserves input names for ImageClassifier and CausalLM models. Refactors test utilities to use SignatureDef for input/output mapping, ensuring meaningful names and robust output verification. Updates numerical accuracy checks to compare outputs by name using SignatureDef, and adds validation for expected input/output names in exported models.
Consolidates LiteRT input preparation, inference, and output verification into clearer helper methods. Improves handling of dynamic shapes, input/output name matching via SignatureDef, and output comparison logic. Updates docstrings and argument names for consistency and readability. bug fix in inference: - there was a bug, that used corrupted the results of invoke, during getting the SignatureDef fixed it.
Updated test utilities to use TFLite's signature runner for inference, simplifying input handling and output extraction. Also updated model creation in numerical accuracy tests to use explicit Input layers for clarity and consistency.
|
/gemini review |
There was a problem hiding this comment.
Code Review
This pull request introduces a robust and well-tested feature for exporting Keras-Hub models to the LiteRT format. The implementation is well-structured, with clear separation of concerns into base classes, model-specific configurations, and the exporter itself. The addition of an export() method to the Task class provides a convenient user-facing API. The necessary compatibility shims for TensorFlow export are well-documented, including warnings about private API usage. The test coverage is excellent, including both unit tests with mocks and end-to-end tests with real model presets, which ensures the reliability of this new feature. My feedback includes a few minor suggestions to improve code clarity and maintainability in the API and task modules.
| `dict`. Dictionary mapping input names to their specifications | ||
| """ | ||
| return { | ||
| "token_ids": keras.layers.InputSpec( |
There was a problem hiding this comment.
how would this work for bert? which also has a sequence id input? or fnet which does not take a padding mask IIRC
we can't in general assume all models have the same inputs after preprocessing. the specifics will vary per model. i wonder if our design is not sufficiently general here
keras_hub/src/export/configs.py
Outdated
| """Exporter configuration for Causal Language Models (GPT, LLaMA, etc.).""" | ||
|
|
||
| MODEL_TYPE = "causal_lm" | ||
| EXPECTED_INPUTS = ["token_ids", "padding_mask"] |
There was a problem hiding this comment.
+1, in general we cannot assume the backbone inputs from task type alone. different models will have different inputs after preprocessing
Now supporting gemma3 multimodal and other models too.
Introduces `test_litert_export` methods to numerous model test files, enabling automated testing of LiteRT export functionality. Also updates SAMPromptEncoder mask handling for LiteRT compatibility and adjusts test input shapes for SAMImageSegmenter. These changes improve coverage and reliability of LiteRT export across text, image, audio, and multimodal models.
Added new exporter configs for AudioToText, SAMImageSegmenter, and TextToImage models. Refactored input signature logic for multimodal and text/image models in configs.py and litert.py. Updated test files to remove unused keras imports and improve input data handling for LiteRT export compatibility.
|
@gemini-code-assist review the complete code |
There was a problem hiding this comment.
Code Review
This pull request introduces a significant new feature: exporting Keras-Hub models to the LiteRT format. The changes are extensive and well-structured, including new base classes for exporters, model-specific configurations, and a comprehensive testing framework. The Task and Backbone classes are updated to support the new export functionality. A large number of tests have been added to ensure the correctness of the export process for various models, which is excellent.
My review has identified a couple of areas for improvement. There's a bug in the ObjectDetectorExporterConfig where it doesn't handle the image_size parameter correctly, and it seems some tests in sam_image_segmenter_test.py might have been removed unintentionally.
Overall, this is a great contribution that significantly enhances the deployment capabilities of Keras-Hub models. Addressing the identified issues will make it even more robust.
keras_hub/src/export/configs.py
Outdated
| def get_input_signature(self, image_size=None): | ||
| """Get input signature for object detector models. | ||
| Args: | ||
| image_size: `int`, `tuple` or `None`. Optional image size. | ||
| Returns: | ||
| `dict`. Dictionary mapping input names to their specifications | ||
| """ | ||
| # Object detectors use dynamic image shapes to support variable input | ||
| # sizes | ||
| # The preprocessor image_size is used for training but export allows any | ||
| # size | ||
| dtype = _infer_image_dtype(self.model) | ||
|
|
||
| return { | ||
| "images": keras.layers.InputSpec( | ||
| dtype=dtype, shape=(None, None, None, 3) | ||
| ), | ||
| "image_shape": keras.layers.InputSpec( | ||
| dtype="int32", shape=(None, 2) | ||
| ), | ||
| } |
There was a problem hiding this comment.
The get_input_signature method for ObjectDetectorExporterConfig currently ignores the image_size parameter and always returns a fully dynamic shape for images. While dynamic shapes are useful, the exporter should also support exporting with a fixed image_size when provided. This would make it consistent with other image-based configurations and is expected by the testing framework.
The method should be updated to use the image_size parameter when it's provided, while still defaulting to dynamic shapes if it's None.
def get_input_signature(self, image_size=None):
"""Get input signature for object detector models.
Args:
image_size: `int`, `tuple` or `None`. Optional image size.
Returns:
`dict`. Dictionary mapping input names to their specifications
"""
dtype = _infer_image_dtype(self.model)
height, width = None, None
if image_size:
if isinstance(image_size, int):
height, width = image_size, image_size
else:
height, width = image_size
return {
"images": keras.layers.InputSpec(
dtype=dtype, shape=(None, height, width, 3)
),
"image_shape": keras.layers.InputSpec(
dtype="int32", shape=(None, 2)
),
}
There was a problem hiding this comment.
Code Review
This pull request introduces a comprehensive model export functionality to LiteRT, which is a significant and valuable addition. The implementation is well-structured, featuring base classes for exporters, specific configurations for various model types, and a clean export() API on the Task model. The inclusion of extensive tests, particularly the run_litert_export_test helper and broad model coverage, is commendable. My review focuses on a minor logic improvement in a test helper and ensuring consistency in applying test decorators for long-running tests.
Introduces DepthEstimatorExporterConfig for depth estimation model export, updates LiteRT exporter to support DepthEstimator, and refines input signature logic for object detectors. Test files now skip LiteRT export tests for non-TensorFlow backends, and output thresholds/statistical comparison modes are adjusted for several models. The base TestCase class is refactored to remove TensorFlow dependency, improving backend compatibility. The run_litert_tests.py script is enhanced for backend selection, parametrized test parsing, and markdown reporting.
|
@gemini-code-assist review |
There was a problem hiding this comment.
Code Review
This pull request introduces a comprehensive model export functionality to the LiteRT format. It adds a new export module with a well-designed set of base classes, model-specific configurations, and a LiteRTExporter. A user-friendly export() method is also added to the Task base class.
The implementation is robust, featuring an adapter mechanism to handle the dictionary-based inputs of Keras-Hub models, making them compatible with the underlying Keras LiteRT exporter. The support for dynamic input shapes is a great feature that enhances flexibility.
The PR is extensive and includes a very thorough suite of tests, covering not only the new export modules but also adding export tests for a wide range of existing models. This significantly increases confidence in the new feature.
The code is well-structured, documented, and adheres to the repository's style guide. I've found one minor issue in a new test helper function that needs to be addressed. Overall, this is an excellent contribution.
Marked LiteRT export tests with @pytest.mark.large in object detector, causal LM, and image segmenter test files to better categorize resource-intensive tests. Also simplified dtype handling in TestCase for TensorFlow tensors.
Test cases for DepthAnything, FNet, PaliGemma, RetinaNet, and SAM models now use statistical comparison modes and output thresholds for LiteRT export tests. The input size for RetinaNet tests was increased to 800. The TestCase utility was improved to extract comparison_mode from export_kwargs and to allow SignatureDef to have additional optional inputs, only failing if expected inputs are missing.
Moved the Dictionary handling to keras Core, Also moved Wrapper creation and handling to keras. Removed Adapter stuff. Simplifies input signatures for image and object detection models to use single tensor inputs instead of dictionaries. Removes support for 'image_shape' as an input for object detectors and updates related tests and exporter logic to match. Cleans up legacy model building code and adapts tests to new input formats for classifiers, detectors, and segmenters. Delegates LiteRT export to Keras Core's exporter for improved input handling.
|
@gemini-code-assist review |
There was a problem hiding this comment.
Code Review
This pull request introduces a significant and valuable feature: exporting Keras-Hub models to LiteRT format. The implementation is well-structured, following a clear pattern with base classes, specific configurations for different model types, and a dedicated exporter. The addition of an export() method to the Task class provides a user-friendly API.
The test coverage is excellent. The new functionality is validated with a comprehensive set of unit tests for the export framework and end-to-end tests for a wide range of production models. This ensures the feature is robust and reliable.
I've found a couple of minor issues in the test helper functions that could be improved for robustness and clarity. Overall, this is a very solid contribution.
| def convert_for_tflite(x): | ||
| """Convert tensor/array to TFLite-compatible dtypes.""" | ||
| if hasattr(x, "dtype"): | ||
| if isinstance(x, np.ndarray): | ||
| if x.dtype == bool: | ||
| return x.astype(np.int32) | ||
| elif x.dtype == np.float64: | ||
| return x.astype(np.float32) | ||
| elif x.dtype == np.int64: | ||
| return x.astype(np.int32) | ||
| else: # TensorFlow tensor | ||
| if x.dtype == tf.bool: | ||
| return tf.cast(x, tf.int32).numpy() | ||
| elif x.dtype == tf.float64: | ||
| return tf.cast(x, tf.float32).numpy() | ||
| elif x.dtype == tf.int64: | ||
| return tf.cast(x, tf.int32).numpy() | ||
| else: | ||
| return x.numpy() if hasattr(x, "numpy") else x | ||
| elif hasattr(x, "numpy"): | ||
| return x.numpy() | ||
| return x | ||
|
|
There was a problem hiding this comment.
The convert_for_tflite function has a bug where it implicitly returns None for np.ndarray inputs if their dtype doesn't match one of the explicit checks (e.g., bool, np.float64, np.int64). This will cause errors for numpy arrays with standard dtypes like np.float32 or np.int32.
You should add a return x for np.ndarray to handle cases where no dtype conversion is needed. You can also simplify the TensorFlow tensor branch by removing a redundant check.
def convert_for_tflite(x):
"""Convert tensor/array to TFLite-compatible dtypes."""
if hasattr(x, "dtype"):
if isinstance(x, np.ndarray):
if x.dtype == bool:
return x.astype(np.int32)
elif x.dtype == np.float64:
return x.astype(np.float32)
elif x.dtype == np.int64:
return x.astype(np.int32)
return x
else: # TensorFlow tensor
if x.dtype == tf.bool:
return tf.cast(x, tf.int32).numpy()
elif x.dtype == tf.float64:
return tf.cast(x, tf.float32).numpy()
elif x.dtype == tf.int64:
return tf.cast(x, tf.int32).numpy()
else:
return x.numpy()
elif hasattr(x, "numpy"):
return x.numpy()
return xRemoved LiteRTExporter class and its tests, updating all usages to call Keras Core's export_litert directly with domain-specific input signatures from exporter configs. The base exporter classes were moved into configs.py, and input signature logic now defaults to preprocessor's sequence_length if available. Updated documentation and tests to reflect the new export flow.
Marked LiteRT export tests as large in multiple model test files to improve test categorization. Also updated statistical comparison parameters in PARSeq test and clarified docstring in Task class. Refactored single tensor input handling in TestCase for consistency.
Applied the @pytest.mark.large decorator to the test_saved_model methods in d_fine_object_detector_test.py, retinanet_object_detector_test.py, and sam_image_segmenter_test.py to categorize these tests as large.
Deleted the _get_save_spec and _trackable_children methods from the Backbone class, which provided compatibility for TensorFlow SavedModel/TFLite export and filtered problematic _DictWrapper objects. Moved this to core keras.
3 participants
This pull request introduces a new export API for Keras-Hub models, providing domain-specific export configurations and a convenience function for exporting models to LiteRT format. It also adds comprehensive tests for the new export configuration classes to ensure correct input signature handling and model type detection.
Export API additions:
keras_hub/api/export/__init__.pyto expose domain-specific exporter configuration classes for various model types (e.g.,CausalLMExporterConfig,ImageClassifierExporterConfig, etc.) in the public API.keras_hub/api/__init__.pyto include the newexportmodule in the public API.Export implementation:
keras_hub/src/export/litert.pywith a convenience functionexport_litertthat automatically detects model type and exports using the correct configuration, leveraging flexible input signatures for Keras-Hub models.keras_hub/src/export/__init__.pyto import all relevant exporter configuration classes and the newexport_litertfunction for internal use.Testing:
keras_hub/src/export/configs_test.pywith extensive tests for all major exporter configuration classes, verifying correct input signatures and expected model types for each domain.Tested models:
TextClassifier Models (8/8 passed):
Unknown Models (10/15 passed):
Failed:
ImageClassifier Models (12/12 passed):
Passed:
ObjectDetector Models (2/2 passed):
Passed:
CausalLM Models (12/13 passed):
Passed:
Failed: