Model Export to liteRT

keras_hub/src/models/albert/albert_text_classifier_test.py

@@ -61,6 +61,14 @@ def test_saved_model(self):
             input_data=self.input_data,
         )
 
+    @pytest.mark.large
+    def test_litert_export(self):
+        self.run_litert_export_test(
+            cls=AlbertTextClassifier,
+            init_kwargs=self.init_kwargs,
+            input_data=self.input_data,
+        )
+
     @pytest.mark.extra_large
     def test_all_presets(self):
         for preset in AlbertTextClassifier.presets:

keras_hub/src/models/bart/bart_seq_2_seq_lm_test.py

@@ -149,6 +149,14 @@ def test_saved_model(self):
             input_data=self.input_data,
         )
 
+    @pytest.mark.large
+    def test_litert_export(self):
+        self.run_litert_export_test(
+            cls=BartSeq2SeqLM,
+            init_kwargs=self.init_kwargs,
+            input_data=self.input_data,
+        )
+
     @pytest.mark.extra_large
     def test_all_presets(self):
         for preset in BartSeq2SeqLM.presets:

keras_hub/src/models/basnet/basnet_test.py

@@ -3,6 +3,9 @@
 
 from keras_hub.src.models.basnet.basnet import BASNetImageSegmenter
 from keras_hub.src.models.basnet.basnet_backbone import BASNetBackbone
+from keras_hub.src.models.basnet.basnet_image_converter import (
+    BASNetImageConverter,
+)
 from keras_hub.src.models.basnet.basnet_preprocessor import BASNetPreprocessor
 from keras_hub.src.models.resnet.resnet_backbone import ResNetBackbone
 from keras_hub.src.tests.test_case import TestCase
@@ -26,7 +29,9 @@ def setUp(self):
             image_encoder=self.image_encoder,
             num_classes=1,
         )
-        self.preprocessor = BASNetPreprocessor()
+        self.preprocessor = BASNetPreprocessor(
+            image_converter=BASNetImageConverter(height=64, width=64)
+        )
         self.init_kwargs = {
             "backbone": self.backbone,
             "preprocessor": self.preprocessor,
@@ -49,6 +54,14 @@ def test_saved_model(self):
             input_data=self.images,
         )
 
+    @pytest.mark.large
+    def test_litert_export(self):
+        self.run_litert_export_test(
+            cls=BASNetImageSegmenter,
+            init_kwargs=self.init_kwargs,
+            input_data=self.images,
+        )
+
     def test_end_to_end_model_predict(self):
         model = BASNetImageSegmenter(**self.init_kwargs)
         output = model.predict(self.images)

keras_hub/src/models/bert/bert_text_classifier_test.py

@@ -53,6 +53,14 @@ def test_saved_model(self):
             input_data=self.input_data,
         )
 
+    @pytest.mark.large
+    def test_litert_export(self):
+        self.run_litert_export_test(
+            cls=BertTextClassifier,
+            init_kwargs=self.init_kwargs,
+            input_data=self.input_data,
+        )
+
     @pytest.mark.extra_large
     def test_smallest_preset(self):
         self.run_preset_test(

keras_hub/src/models/bloom/bloom_causal_lm_test.py

@@ -164,6 +164,14 @@ def test_saved_model(self):
             input_data=self.input_data,
         )
 
+    @pytest.mark.large
+    def test_litert_export(self):
+        self.run_litert_export_test(
+            cls=BloomCausalLM,
+            init_kwargs=self.init_kwargs,
+            input_data=self.input_data,
+        )
+
     @pytest.mark.extra_large
     def test_all_presets(self):
         for preset in BloomCausalLM.presets:

keras_hub/src/models/clip/clip_backbone_test.py

@@ -44,6 +44,16 @@ def test_saved_model(self):
             input_data=self.input_data,
         )
 
+    @pytest.mark.large
+    def test_litert_export(self):
+        self.run_litert_export_test(
+            cls=CLIPBackbone,
+            init_kwargs=self.init_kwargs,
+            input_data=self.input_data,
+            comparison_mode="statistical",
+            output_thresholds={"*": {"max": 1e-4, "mean": 1e-5}},
+        )
+
     @pytest.mark.extra_large
     def test_all_presets(self):
         for preset in CLIPBackbone.presets:

keras_hub/src/models/cspnet/cspnet_image_classifier_test.py

@@ -76,3 +76,11 @@ def test_saved_model(self):
             init_kwargs=self.init_kwargs,
             input_data=self.images,
         )
+
+    @pytest.mark.large
+    def test_litert_export(self):
+        self.run_litert_export_test(
+            cls=CSPNetImageClassifier,
+            init_kwargs=self.init_kwargs,
+            input_data=self.images,
+        )

keras_hub/src/models/d_fine/d_fine_object_detector_test.py

@@ -152,3 +152,30 @@ def test_saved_model(self):
             init_kwargs=init_kwargs,
             input_data=self.images,
         )
+
+    @pytest.mark.large
+    def test_litert_export(self):
+        backbone = DFineBackbone(**self.base_backbone_kwargs)
+        init_kwargs = {
+            "backbone": backbone,
+            "num_classes": 4,
+            "bounding_box_format": self.bounding_box_format,
+            "preprocessor": self.preprocessor,
+        }
+
+        # D-Fine ObjectDetector only takes images as input
+        input_data = self.images
+
+        self.run_litert_export_test(
+            cls=DFineObjectDetector,
+            init_kwargs=init_kwargs,
+            input_data=input_data,
+            comparison_mode="statistical",
+            output_thresholds={
+                "intermediate_predicted_corners": {"max": 5.0, "mean": 0.05},
+                "intermediate_logits": {"max": 5.0, "mean": 0.1},
+                "enc_topk_logits": {"max": 5.0, "mean": 0.03},
+                "logits": {"max": 2.0, "mean": 0.03},
+                "*": {"max": 1.0, "mean": 0.03},
+            },
+        )

keras_hub/src/models/deberta_v3/deberta_v3_text_classifier_test.py

@@ -64,6 +64,14 @@ def test_saved_model(self):
             input_data=self.input_data,
         )
 
+    @pytest.mark.large
+    def test_litert_export(self):
+        self.run_litert_export_test(
+            cls=DebertaV3TextClassifier,
+            init_kwargs=self.init_kwargs,
+            input_data=self.input_data,
+        )
+
     @pytest.mark.extra_large
     def test_all_presets(self):
         for preset in DebertaV3TextClassifier.presets:

keras_hub/src/models/deeplab_v3/deeplab_v3_segmenter_test.py

@@ -70,3 +70,15 @@ def test_saved_model(self):
             init_kwargs=self.init_kwargs,
             input_data=self.images,
         )
+
+    @pytest.mark.large
+    def test_litert_export(self):
+        self.run_litert_export_test(
+            cls=DeepLabV3ImageSegmenter,
+            init_kwargs=self.init_kwargs,
+            input_data=self.images,
+            comparison_mode="statistical",
+            output_thresholds={
+                "*": {"max": 0.6, "mean": 0.3},
+            },
+        )

keras_hub/src/models/deit/deit_image_classifier_test.py

@@ -55,3 +55,11 @@ def test_saved_model(self):
             init_kwargs=self.init_kwargs,
             input_data=self.images,
         )
+
+    @pytest.mark.large
+    def test_litert_export(self):
+        self.run_litert_export_test(
+            cls=DeiTImageClassifier,
+            init_kwargs=self.init_kwargs,
+            input_data=self.images,
+        )

keras_hub/src/models/densenet/densenet_image_classifier_test.py

@@ -61,6 +61,14 @@ def test_saved_model(self):
             input_data=self.images,
         )
 
+    @pytest.mark.large
+    def test_litert_export(self):
+        self.run_litert_export_test(
+            cls=DenseNetImageClassifier,
+            init_kwargs=self.init_kwargs,
+            input_data=self.images,
+        )
+
     @pytest.mark.extra_large
     def test_all_presets(self):
         for preset in DenseNetImageClassifier.presets:

keras_hub/src/models/depth_anything/depth_anything_depth_estimator_test.py

@@ -85,6 +85,16 @@ def test_saved_model(self):
             input_data=self.images,
         )
 
+    @pytest.mark.large
+    def test_litert_export(self):
+        self.run_litert_export_test(
+            cls=DepthAnythingDepthEstimator,
+            init_kwargs=self.init_kwargs,
+            input_data=self.images,
+            comparison_mode="statistical",
+            output_thresholds={"depths": {"max": 2e-4, "mean": 1e-5}},
+        )
+
     @pytest.mark.extra_large
     def test_all_presets(self):
         images = np.ones((2, 518, 518, 3), dtype="float32")

keras_hub/src/models/dinov2/dinov2_backbone_test.py

@@ -49,6 +49,16 @@ def test_saved_model(self):
             input_data=self.input_data,
         )
 
+    @pytest.mark.large
+    def test_litert_export(self):
+        self.run_litert_export_test(
+            cls=DINOV2Backbone,
+            init_kwargs=self.init_kwargs,
+            input_data=self.input_data,
+            comparison_mode="statistical",
+            output_thresholds={"*": {"max": 1e-4, "mean": 1e-5}},
+        )
+
     @pytest.mark.large
     def test_position_embedding_interpolation(self):
         model = DINOV2Backbone(**self.init_kwargs)
@@ -144,6 +154,16 @@ def test_saved_model(self):
             input_data=self.input_data,
         )
 
+    @pytest.mark.large
+    def test_litert_export(self):
+        self.run_litert_export_test(
+            cls=DINOV2Backbone,
+            init_kwargs=self.init_kwargs,
+            input_data=self.input_data,
+            comparison_mode="statistical",
+            output_thresholds={"*": {"max": 1e-4, "mean": 1e-5}},
+        )
+
     @pytest.mark.kaggle_key_required
     @pytest.mark.extra_large
     def test_smallest_preset(self):

keras_hub/src/models/dinov3/dinov3_backbone_test.py

@@ -52,6 +52,14 @@ def test_saved_model(self):
             input_data=self.input_data,
         )
 
+    @pytest.mark.large
+    def test_litert_export(self):
+        self.run_litert_export_test(
+            cls=DINOV3Backbone,
+            init_kwargs=self.init_kwargs,
+            input_data=self.input_data,
+        )
+
     @pytest.mark.large
     def test_position_embedding_interpolation(self):
         model = DINOV3Backbone(**self.init_kwargs)

keras_hub/src/models/distil_bert/distil_bert_text_classifier_test.py

@@ -59,6 +59,14 @@ def test_saved_model(self):
             input_data=self.input_data,
         )
 
+    @pytest.mark.large
+    def test_litert_export(self):
+        self.run_litert_export_test(
+            cls=DistilBertTextClassifier,
+            init_kwargs=self.init_kwargs,
+            input_data=self.input_data,
+        )
+
     @pytest.mark.extra_large
     def test_all_presets(self):
         for preset in DistilBertTextClassifier.presets:

keras_hub/src/models/efficientnet/efficientnet_image_classifier_test.py

@@ -7,6 +7,12 @@
 from keras_hub.src.models.efficientnet.efficientnet_image_classifier import (
     EfficientNetImageClassifier,
 )
+from keras_hub.src.models.efficientnet.efficientnet_image_classifier_preprocessor import (  # noqa: E501
+    EfficientNetImageClassifierPreprocessor,
+)
+from keras_hub.src.models.efficientnet.efficientnet_image_converter import (
+    EfficientNetImageConverter,
+)
 from keras_hub.src.tests.test_case import TestCase
 
 
@@ -38,6 +44,9 @@ def setUp(self):
         self.init_kwargs = {
             "backbone": backbone,
             "num_classes": 1000,
+            "preprocessor": EfficientNetImageClassifierPreprocessor(
+                image_converter=EfficientNetImageConverter(image_size=(16, 16))
+            ),
         }
         self.train_data = (self.images, self.labels)
 
@@ -82,3 +91,11 @@ def test_all_presets(self):
                 input_data=self.images,
                 expected_output_shape=(2, 2),
             )
+
+    @pytest.mark.large
+    def test_litert_export(self):
+        self.run_litert_export_test(
+            cls=EfficientNetImageClassifier,
+            init_kwargs=self.init_kwargs,
+            input_data=self.images,
+        )

keras_hub/src/models/electra/electra_backbone_test.py

@@ -41,6 +41,14 @@ def test_saved_model(self):
             input_data=self.input_data,
         )
 
+    @pytest.mark.large
+    def test_litert_export(self):
+        self.run_litert_export_test(
+            cls=ElectraBackbone,
+            init_kwargs=self.init_kwargs,
+            input_data=self.input_data,
+        )
+
     @pytest.mark.extra_large
     def test_smallest_preset(self):
         self.run_preset_test(

keras_hub/src/models/esm/esm_classifier_test.py

@@ -51,3 +51,11 @@ def test_saved_model(self):
             init_kwargs=self.init_kwargs,
             input_data=self.input_data,
         )
+
+    @pytest.mark.large
+    def test_litert_export(self):
+        self.run_litert_export_test(
+            cls=ESMProteinClassifier,
+            init_kwargs=self.init_kwargs,
+            input_data=self.input_data,
+        )

keras_hub/src/models/f_net/f_net_text_classifier_test.py

@@ -57,6 +57,20 @@ def test_saved_model(self):
             input_data=self.input_data,
         )
 
+    @pytest.mark.large
+    def test_litert_export(self):
+        # F-Net does NOT use padding_mask - it only uses token_ids and
+        # segment_ids. Don't add padding_mask to input_data.
+        self.run_litert_export_test(
+            cls=FNetTextClassifier,
+            init_kwargs=self.init_kwargs,
+            input_data=self.input_data,
+            comparison_mode="statistical",
+            output_thresholds={
+                "*": {"max": 0.01, "mean": 0.005},
+            },
+        )
+
     @pytest.mark.extra_large
     def test_all_presets(self):
         for preset in FNetTextClassifier.presets:

keras_hub/src/models/falcon/falcon_causal_lm_test.py

@@ -164,6 +164,14 @@ def test_saved_model(self):
             input_data=self.input_data,
         )
 
+    @pytest.mark.large
+    def test_litert_export(self):
+        self.run_litert_export_test(
+            cls=FalconCausalLM,
+            init_kwargs=self.init_kwargs,
+            input_data=self.input_data,
+        )
+
     @pytest.mark.extra_large
     def test_all_presets(self):
         for preset in FalconCausalLM.presets: