diff --git a/keras_hub/api/layers/__init__.py b/keras_hub/api/layers/__init__.py index a32da3a42a..9a310ad227 100644 --- a/keras_hub/api/layers/__init__.py +++ b/keras_hub/api/layers/__init__.py @@ -40,9 +40,7 @@ from keras_hub.src.models.densenet.densenet_image_converter import ( DenseNetImageConverter, ) -from keras_hub.src.models.mix_transformer.mix_transformer_image_converter import ( - MiTImageConverter, -) +from keras_hub.src.models.mit.mit_image_converter import MiTImageConverter from keras_hub.src.models.pali_gemma.pali_gemma_image_converter import ( PaliGemmaImageConverter, ) diff --git a/keras_hub/api/models/__init__.py b/keras_hub/api/models/__init__.py index 4f16fcc3eb..408327f9cc 100644 --- a/keras_hub/api/models/__init__.py +++ b/keras_hub/api/models/__init__.py @@ -202,13 +202,9 @@ MistralCausalLMPreprocessor, ) from keras_hub.src.models.mistral.mistral_tokenizer import MistralTokenizer -from keras_hub.src.models.mix_transformer.mix_transformer_backbone import ( - MiTBackbone, -) -from keras_hub.src.models.mix_transformer.mix_transformer_classifier import ( - MiTImageClassifier, -) -from keras_hub.src.models.mix_transformer.mix_transformer_classifier_preprocessor import ( +from keras_hub.src.models.mit.mit_backbone import MiTBackbone +from keras_hub.src.models.mit.mit_image_classifier import MiTImageClassifier +from keras_hub.src.models.mit.mit_image_classifier_preprocessor import ( MiTImageClassifierPreprocessor, ) from keras_hub.src.models.mobilenet.mobilenet_backbone import MobileNetBackbone diff --git a/keras_hub/src/models/mit/__init__.py b/keras_hub/src/models/mit/__init__.py new file mode 100644 index 0000000000..f581202b1c --- /dev/null +++ b/keras_hub/src/models/mit/__init__.py @@ -0,0 +1,6 @@ +from keras_hub.src.models.mit.mit_backbone import MiTBackbone +from keras_hub.src.models.mit.mit_image_classifier import MiTImageClassifier +from keras_hub.src.models.mit.mit_presets import backbone_presets +from keras_hub.src.utils.preset_utils import register_presets + +register_presets(backbone_presets, MiTBackbone) diff --git a/keras_hub/src/models/mix_transformer/mix_transformer_backbone.py b/keras_hub/src/models/mit/mit_backbone.py similarity index 96% rename from keras_hub/src/models/mix_transformer/mix_transformer_backbone.py rename to keras_hub/src/models/mit/mit_backbone.py index e8f881aee3..1398fe8eee 100644 --- a/keras_hub/src/models/mix_transformer/mix_transformer_backbone.py +++ b/keras_hub/src/models/mit/mit_backbone.py @@ -4,12 +4,8 @@ from keras_hub.src.api_export import keras_hub_export from keras_hub.src.models.feature_pyramid_backbone import FeaturePyramidBackbone -from keras_hub.src.models.mix_transformer.mix_transformer_layers import ( - HierarchicalTransformerEncoder, -) -from keras_hub.src.models.mix_transformer.mix_transformer_layers import ( - OverlappingPatchingAndEmbedding, -) +from keras_hub.src.models.mit.mit_layers import HierarchicalTransformerEncoder +from keras_hub.src.models.mit.mit_layers import OverlappingPatchingAndEmbedding @keras_hub_export("keras_hub.models.MiTBackbone") diff --git a/keras_hub/src/models/mit/mit_backbone_test.py b/keras_hub/src/models/mit/mit_backbone_test.py new file mode 100644 index 0000000000..553a266e5b --- /dev/null +++ b/keras_hub/src/models/mit/mit_backbone_test.py @@ -0,0 +1,45 @@ +import numpy as np +import pytest + +from keras_hub.src.models.mit.mit_backbone import MiTBackbone +from keras_hub.src.tests.test_case import TestCase + + +class MiTBackboneTest(TestCase): + def setUp(self): + self.init_kwargs = { + "depths": [2, 2], + "image_shape": (16, 16, 3), + "hidden_dims": [4, 8], + "num_layers": 2, + "blockwise_num_heads": [1, 2], + "blockwise_sr_ratios": [8, 4], + "max_drop_path_rate": 0.1, + "patch_sizes": [7, 3], + "strides": [4, 2], + } + self.input_size = 16 + self.input_data = np.ones( + (2, self.input_size, self.input_size, 3), dtype="float32" + ) + + def test_backbone_basics(self): + self.run_vision_backbone_test( + cls=MiTBackbone, + init_kwargs=self.init_kwargs, + input_data=self.input_data, + expected_output_shape=(2, 2, 2, 8), + expected_pyramid_output_keys=["P1", "P2"], + expected_pyramid_image_sizes=[(4, 4), (2, 2)], + run_quantization_check=False, + run_mixed_precision_check=False, + run_data_format_check=False, + ) + + @pytest.mark.large + def test_saved_model(self): + self.run_model_saving_test( + cls=MiTBackbone, + init_kwargs=self.init_kwargs, + input_data=self.input_data, + ) diff --git a/keras_hub/src/models/mix_transformer/mix_transformer_classifier.py b/keras_hub/src/models/mit/mit_image_classifier.py similarity index 64% rename from keras_hub/src/models/mix_transformer/mix_transformer_classifier.py rename to keras_hub/src/models/mit/mit_image_classifier.py index beab6646ba..370920ddf9 100644 --- a/keras_hub/src/models/mix_transformer/mix_transformer_classifier.py +++ b/keras_hub/src/models/mit/mit_image_classifier.py @@ -1,9 +1,7 @@ from keras_hub.src.api_export import keras_hub_export from keras_hub.src.models.image_classifier import ImageClassifier -from keras_hub.src.models.mix_transformer.mix_transformer_backbone import ( - MiTBackbone, -) -from keras_hub.src.models.mix_transformer.mix_transformer_classifier_preprocessor import ( +from keras_hub.src.models.mit.mit_backbone import MiTBackbone +from keras_hub.src.models.mit.mit_image_classifier_preprocessor import ( MiTImageClassifierPreprocessor, ) diff --git a/keras_hub/src/models/mix_transformer/mix_transformer_classifier_preprocessor.py b/keras_hub/src/models/mit/mit_image_classifier_preprocessor.py similarity index 64% rename from keras_hub/src/models/mix_transformer/mix_transformer_classifier_preprocessor.py rename to keras_hub/src/models/mit/mit_image_classifier_preprocessor.py index 61c994c5fb..d3859c30d6 100644 --- a/keras_hub/src/models/mix_transformer/mix_transformer_classifier_preprocessor.py +++ b/keras_hub/src/models/mit/mit_image_classifier_preprocessor.py @@ -2,12 +2,8 @@ from keras_hub.src.models.image_classifier_preprocessor import ( ImageClassifierPreprocessor, ) -from keras_hub.src.models.mix_transformer.mix_transformer_backbone import ( - MiTBackbone, -) -from keras_hub.src.models.mix_transformer.mix_transformer_image_converter import ( - MiTImageConverter, -) +from keras_hub.src.models.mit.mit_backbone import MiTBackbone +from keras_hub.src.models.mit.mit_image_converter import MiTImageConverter @keras_hub_export("keras_hub.models.MiTImageClassifierPreprocessor") diff --git a/keras_hub/src/models/mix_transformer/mix_transformer_classifier_test.py b/keras_hub/src/models/mit/mit_image_classifier_test.py similarity index 87% rename from keras_hub/src/models/mix_transformer/mix_transformer_classifier_test.py rename to keras_hub/src/models/mit/mit_image_classifier_test.py index fb7ff5ce2b..058ea4fbb2 100644 --- a/keras_hub/src/models/mix_transformer/mix_transformer_classifier_test.py +++ b/keras_hub/src/models/mit/mit_image_classifier_test.py @@ -1,12 +1,8 @@ import numpy as np import pytest -from keras_hub.src.models.mix_transformer.mix_transformer_backbone import ( - MiTBackbone, -) -from keras_hub.src.models.mix_transformer.mix_transformer_classifier import ( - MiTImageClassifier, -) +from keras_hub.src.models.mit.mit_backbone import MiTBackbone +from keras_hub.src.models.mit.mit_image_classifier import MiTImageClassifier from keras_hub.src.tests.test_case import TestCase diff --git a/keras_hub/src/models/mix_transformer/mix_transformer_image_converter.py b/keras_hub/src/models/mit/mit_image_converter.py similarity index 81% rename from keras_hub/src/models/mix_transformer/mix_transformer_image_converter.py rename to keras_hub/src/models/mit/mit_image_converter.py index e59ea26b66..269fcb88fd 100644 --- a/keras_hub/src/models/mix_transformer/mix_transformer_image_converter.py +++ b/keras_hub/src/models/mit/mit_image_converter.py @@ -1,6 +1,6 @@ from keras_hub.src.api_export import keras_hub_export from keras_hub.src.layers.preprocessing.image_converter import ImageConverter -from keras_hub.src.models.mix_transformer import MiTBackbone +from keras_hub.src.models.mit import MiTBackbone @keras_hub_export("keras_hub.layers.MiTImageConverter") diff --git a/keras_hub/src/models/mix_transformer/mix_transformer_layers.py b/keras_hub/src/models/mit/mit_layers.py similarity index 100% rename from keras_hub/src/models/mix_transformer/mix_transformer_layers.py rename to keras_hub/src/models/mit/mit_layers.py diff --git a/keras_hub/src/models/mix_transformer/mix_transformer_presets.py b/keras_hub/src/models/mit/mit_presets.py similarity index 100% rename from keras_hub/src/models/mix_transformer/mix_transformer_presets.py rename to keras_hub/src/models/mit/mit_presets.py diff --git a/keras_hub/src/models/mix_transformer/mix_transformer_backbone_test.py b/keras_hub/src/models/mit/mix_transformer_backbone_test.py similarity index 93% rename from keras_hub/src/models/mix_transformer/mix_transformer_backbone_test.py rename to keras_hub/src/models/mit/mix_transformer_backbone_test.py index b3840f5c07..553a266e5b 100644 --- a/keras_hub/src/models/mix_transformer/mix_transformer_backbone_test.py +++ b/keras_hub/src/models/mit/mix_transformer_backbone_test.py @@ -1,9 +1,7 @@ import numpy as np import pytest -from keras_hub.src.models.mix_transformer.mix_transformer_backbone import ( - MiTBackbone, -) +from keras_hub.src.models.mit.mit_backbone import MiTBackbone from keras_hub.src.tests.test_case import TestCase diff --git a/keras_hub/src/models/mix_transformer/__init__.py b/keras_hub/src/models/mix_transformer/__init__.py deleted file mode 100644 index f2292a35ad..0000000000 --- a/keras_hub/src/models/mix_transformer/__init__.py +++ /dev/null @@ -1,12 +0,0 @@ -from keras_hub.src.models.mix_transformer.mix_transformer_backbone import ( - MiTBackbone, -) -from keras_hub.src.models.mix_transformer.mix_transformer_classifier import ( - MiTImageClassifier, -) -from keras_hub.src.models.mix_transformer.mix_transformer_presets import ( - backbone_presets, -) -from keras_hub.src.utils.preset_utils import register_presets - -register_presets(backbone_presets, MiTBackbone)