keras-hub-nightly 0.16.1.dev202410020340__py3-none-any.whl → 0.19.0.dev202501260345__py3-none-any.whl

keras_hub/src/models/llama/llama_causal_lm.py

@@ -42,7 +42,9 @@ class LlamaCausalLM(CausalLM):
         self.preprocessor = preprocessor
 
         # === Functional Model ===
-        inputs = backbone.inputs
+        # This must be "backbone.input" i.e. the full input structure,
+        # rather than "backbone.inputs" which is the flattened list of inputs.
+        inputs = backbone.input
         hidden_states = backbone(inputs)
         outputs = backbone.token_embedding(hidden_states, reverse=True)
         super().__init__(

keras_hub/src/models/llama/llama_presets.py

@@ -6,11 +6,9 @@ backbone_presets = {
         "metadata": {
             "description": "7 billion parameter, 32-layer, base LLaMA 2 model.",
             "params": 6738415616,
-            "official_name": "LLaMA 2",
-            "path": "llama2",
-            "model_card": "https://github.com/meta-llama/llama",
+            "path": "llama",
         },
-        "kaggle_handle": "kaggle://keras/llama2/keras/llama2_7b_en/1",
+        "kaggle_handle": "kaggle://keras/llama2/keras/llama2_7b_en/2",
     },
     "llama2_7b_en_int8": {
         "metadata": {
@@ -19,11 +17,9 @@ backbone_presets = {
                 "activation and weights quantized to int8."
             ),
             "params": 6739839488,
-            "official_name": "LLaMA 2",
-            "path": "llama2",
-            "model_card": "https://github.com/meta-llama/llama",
+            "path": "llama",
         },
-        "kaggle_handle": "kaggle://keras/llama2/keras/llama2_7b_en_int8/1",
+        "kaggle_handle": "kaggle://keras/llama2/keras/llama2_7b_en_int8/2",
     },
     "llama2_instruct_7b_en": {
         "metadata": {
@@ -32,11 +28,9 @@ backbone_presets = {
                 "model."
             ),
             "params": 6738415616,
-            "official_name": "LLaMA 2",
-            "path": "llama2",
-            "model_card": "https://github.com/meta-llama/llama",
+            "path": "llama",
         },
-        "kaggle_handle": "kaggle://keras/llama2/keras/llama2_instruct_7b_en/1",
+        "kaggle_handle": "kaggle://keras/llama2/keras/llama2_instruct_7b_en/2",
     },
     "llama2_instruct_7b_en_int8": {
         "metadata": {
@@ -45,11 +39,9 @@ backbone_presets = {
                 "model with activation and weights quantized to int8."
             ),
             "params": 6739839488,
-            "official_name": "LLaMA 2",
-            "path": "llama2",
-            "model_card": "https://github.com/meta-llama/llama",
+            "path": "llama",
         },
-        "kaggle_handle": "kaggle://keras/llama2/keras/llama2_instruct_7b_en_int8/1",
+        "kaggle_handle": "kaggle://keras/llama2/keras/llama2_instruct_7b_en_int8/2",
     },
     "vicuna_1.5_7b_en": {
         "metadata": {
@@ -58,10 +50,8 @@ backbone_presets = {
                 "model."
             ),
             "params": 6738415616,
-            "official_name": "Vicuna",
-            "path": "vicuna",
-            "model_card": "https://github.com/lm-sys/FastChat",
+            "path": "llama",
         },
-        "kaggle_handle": "kaggle://keras/vicuna/keras/vicuna_1.5_7b_en/1",
+        "kaggle_handle": "kaggle://keras/vicuna/keras/vicuna_1.5_7b_en/2",
     },
 }

keras_hub/src/models/llama3/llama3_backbone.py

@@ -24,17 +24,18 @@ class Llama3Backbone(LlamaBackbone):
         num_layers (int): The number of transformer layers.
         num_query_heads (int): The number of query attention heads for
             each transformer.
-        hidden_dim (int): The size of the transformer encoding and pooling layers.
-        intermediate_dim (int): The output dimension of the first Dense layer in a
-            three-layer feedforward network for each transformer.
-        num_key_value_heads (int): The number of key and value attention heads for
-            each transformer.
-        rope_max_wavelength (int, optional): The maximum angular wavelength of the
-            sine/cosine curves, for rotary embeddings. Defaults to `10000`.
-        rope_scaling_factor (float, optional): The scaling factor for calculation
-            of roatary embedding. Defaults to `1.0`.
-        layer_norm_epsilon (float, optional): Epsilon for the layer normalization
-            layers in the transformer decoder. Defaults to `1e-6`.
+        hidden_dim (int): The size of the transformer encoding and pooling
+            layers.
+        intermediate_dim (int): The output dimension of the first Dense layer in
+            a three-layer feedforward network for each transformer.
+        num_key_value_heads (int): The number of key and value attention heads
+            fo each transformer.
+        rope_max_wavelength (int, optional): The maximum angular wavelength of
+            the sine/cosine curves, for rotary embeddings. Defaults to `10000`.
+        rope_scaling_factor (float, optional): The scaling factor for
+            calculation of roatary embedding. Defaults to `1.0`.
+        layer_norm_epsilon (float, optional): Epsilon for the layer
+            normalization layers in the transformer decoder. Defaults to `1e-6`.
         dtype: string or `keras.mixed_precision.DTypePolicy`. The dtype to use
             for model computations and weights. Note that some computations,
             such as softmax and layer normalization, will always be done at

keras_hub/src/models/llama3/llama3_causal_lm.py

@@ -1,9 +1,9 @@
 from keras_hub.src.api_export import keras_hub_export
+from keras_hub.src.models.llama.llama_causal_lm import LlamaCausalLM
 from keras_hub.src.models.llama3.llama3_backbone import Llama3Backbone
 from keras_hub.src.models.llama3.llama3_causal_lm_preprocessor import (
     Llama3CausalLMPreprocessor,
 )
-from keras_hub.src.models.llama.llama_causal_lm import LlamaCausalLM
 
 
 @keras_hub_export("keras_hub.models.Llama3CausalLM")

keras_hub/src/models/llama3/llama3_presets.py

@@ -6,11 +6,9 @@ backbone_presets = {
         "metadata": {
             "description": "8 billion parameter, 32-layer, base LLaMA 3 model.",
             "params": 8030261248,
-            "official_name": "LLaMA 3",
             "path": "llama3",
-            "model_card": "https://github.com/meta-llama/llama3",
         },
-        "kaggle_handle": "kaggle://keras/llama3/keras/llama3_8b_en/3",
+        "kaggle_handle": "kaggle://keras/llama3/keras/llama3_8b_en/4",
     },
     "llama3_8b_en_int8": {
         "metadata": {
@@ -19,11 +17,9 @@ backbone_presets = {
                 "activation and weights quantized to int8."
             ),
             "params": 8031894016,
-            "official_name": "LLaMA 3",
             "path": "llama3",
-            "model_card": "https://github.com/meta-llama/llama3",
         },
-        "kaggle_handle": "kaggle://keras/llama3/keras/llama3_8b_en_int8/1",
+        "kaggle_handle": "kaggle://keras/llama3/keras/llama3_8b_en_int8/2",
     },
     "llama3_instruct_8b_en": {
         "metadata": {
@@ -32,11 +28,9 @@ backbone_presets = {
                 "model."
             ),
             "params": 8030261248,
-            "official_name": "LLaMA 3",
             "path": "llama3",
-            "model_card": "https://github.com/meta-llama/llama3",
         },
-        "kaggle_handle": "kaggle://keras/llama3/keras/llama3_instruct_8b_en/3",
+        "kaggle_handle": "kaggle://keras/llama3/keras/llama3_instruct_8b_en/4",
     },
     "llama3_instruct_8b_en_int8": {
         "metadata": {
@@ -45,12 +39,10 @@ backbone_presets = {
                 "model with activation and weights quantized to int8."
             ),
             "params": 8031894016,
-            "official_name": "LLaMA 3",
             "path": "llama3",
-            "model_card": "https://github.com/meta-llama/llama3",
         },
         "kaggle_handle": (
-            "kaggle://keras/llama3/keras/llama3_instruct_8b_en_int8/1"
+            "kaggle://keras/llama3/keras/llama3_instruct_8b_en_int8/2"
         ),
     },
 }

keras_hub/src/models/llama3/llama3_tokenizer.py

@@ -16,10 +16,33 @@ class Llama3Tokenizer(BytePairTokenizer):
         self,
         vocabulary=None,
         merges=None,
+        bos_token="<|begin_of_text|>",
+        eos_token="<|end_of_text|>",
+        misc_special_tokens={"<|start_header_id|>", "<|end_header_id|>"},
         **kwargs,
     ):
-        self._add_special_token("<|begin_of_text|>", "start_token")
-        self._add_special_token("<|end_of_text|>", "end_token")
+        # Note: all special tokens must also appear in "vocabulary"
+
+        self._add_special_token(bos_token, "start_token")
+        misc_special_tokens -= {bos_token}
+        self._add_special_token(eos_token, "end_token")
+        misc_special_tokens -= {eos_token}
+        for i, token in enumerate(misc_special_tokens):
+            self._add_special_token(token, f"special_token_{i:03d}")
+
+        # Hack:
+        # Llama models use the <|end_of_text|> or the <|eot_id|> as the stop
+        # token. This info can be read from config when loading a Hugging Face
+        # checkpoint but no such config exists for Keras checkpoints.
+        # Setting both probable end tokens when no config is availble will
+        # make text generation work in all cases as it will stop
+        # on both end tokens. However, the packer will always use
+        # "<|end_of_text|>" , which will be the wrong eos_token for "instruct"
+        # variants of Llama3.
+        # TODO: load this correctly from a Keras tokenizer config.
+        if eos_token == "<|end_of_text|>":
+            self._add_special_token("<|eot_id|>", "end_token2")
+
         self.pad_token_id = 0
         super().__init__(
             vocabulary=vocabulary,

keras_hub/src/models/mistral/mistral_backbone.py

@@ -38,22 +38,23 @@ class MistralBackbone(Backbone):
         num_layers (int): The number of transformer layers.
         num_query_heads (int): The number of query attention heads for
             each transformer.
-        hidden_dim (int): The size of the transformer encoding and pooling layers.
-        intermediate_dim (int): The output dimension of the first Dense layer in a
-            three-layer feedforward network for each transformer.
-        num_key_value_heads (int): The number of key and value attention heads for
-            each transformer.
-        rope_max_wavelength (int, optional): The maximum angular wavelength of the
-            sine/cosine curves, for rotary embeddings. Defaults to `10000`.
-        rope_scaling_factor (float, optional): The scaling factor for calculation
-            of roatary embedding. Defaults to `1.0`.
-        layer_norm_epsilon (float, optional): Epsilon for the layer normalization
-            layers in the transformer decoder. Defaults to `1e-6`.
+        hidden_dim (int): The size of the transformer encoding and pooling
+            layers.
+        intermediate_dim (int): The output dimension of the first Dense layer
+            in a three-layer feedforward network for each transformer.
+        num_key_value_heads (int): The number of key and value attention heads
+            for each transformer.
+        rope_max_wavelength (int, optional): The maximum angular wavelength of
+            the sine/cosine curves, for rotary embeddings. Defaults to `10000`.
+        rope_scaling_factor (float, optional): The scaling factor for
+            calculation of roatary embedding. Defaults to `1.0`.
+        layer_norm_epsilon (float, optional): Epsilon for the layer
+            normalization layers in the transformer decoder. Defaults to `1e-6`.
         sliding_window (int, optional): The sliding window for the mistral
-            attention layers. This controls the maximum cache size for the attention
-            layers in each transformer decoder. Only `sliding_window` number of tokens
-            are saved in the cache and used to generate the next token.
-            Defaults to `512`.
+            attention layers. This controls the maximum cache size for the
+            attention layers in each transformer decoder. Only `sliding_window`
+            number of tokens are saved in the cache and used to generate the
+            next token. Defaults to `512`.
         dtype: string or `keras.mixed_precision.DTypePolicy`. The dtype to use
             for model computations and weights. Note that some computations,
             such as softmax and layer normalization, will always be done at

keras_hub/src/models/mistral/mistral_causal_lm.py

@@ -28,9 +28,9 @@ class MistralCausalLM(CausalLM):
 
     Args:
         backbone: A `keras_hub.models.MistralBackbone` instance.
-        preprocessor: A `keras_hub.models.MistralCausalLMPreprocessor` or `None`.
-            If `None`, this model will not apply preprocessing, and inputs
-            should be preprocessed before calling the model.
+        preprocessor: A `keras_hub.models.MistralCausalLMPreprocessor` or
+            `None`. If `None`, this model will not apply preprocessing, and
+            inputs should be preprocessed before calling the model.
     """
 
     backbone_cls = MistralBackbone
@@ -42,7 +42,9 @@ class MistralCausalLM(CausalLM):
         self.preprocessor = preprocessor
 
         # === Functional Model ===
-        inputs = backbone.inputs
+        # This must be "backbone.input" i.e. the full input structure,
+        # rather than "backbone.inputs" which is the flattened list of inputs.
+        inputs = backbone.input
         hidden_states = backbone(inputs)
         outputs = backbone.token_embedding(hidden_states, reverse=True)
         super().__init__(

keras_hub/src/models/mistral/mistral_presets.py

@@ -6,30 +6,24 @@ backbone_presets = {
         "metadata": {
             "description": "Mistral 7B base model",
             "params": 7241732096,
-            "official_name": "Mistral",
             "path": "mistral",
-            "model_card": "https://github.com/mistralai/mistral-src/blob/main/README.md",
         },
-        "kaggle_handle": "kaggle://keras/mistral/keras/mistral_7b_en/6",
+        "kaggle_handle": "kaggle://keras/mistral/keras/mistral_7b_en/7",
     },
     "mistral_instruct_7b_en": {
         "metadata": {
             "description": "Mistral 7B instruct model",
             "params": 7241732096,
-            "official_name": "Mistral",
             "path": "mistral",
-            "model_card": "https://github.com/mistralai/mistral-src/blob/main/README.md",
         },
-        "kaggle_handle": "kaggle://keras/mistral/keras/mistral_instruct_7b_en/6",
+        "kaggle_handle": "kaggle://keras/mistral/keras/mistral_instruct_7b_en/7",
     },
     "mistral_0.2_instruct_7b_en": {
         "metadata": {
             "description": "Mistral 7B instruct Version 0.2 model",
             "params": 7241732096,
-            "official_name": "Mistral",
             "path": "mistral",
-            "model_card": "https://github.com/mistralai/mistral-src/blob/main/README.md",
         },
-        "kaggle_handle": "kaggle://keras/mistral/keras/mistral_0.2_instruct_7b_en/1",
+        "kaggle_handle": "kaggle://keras/mistral/keras/mistral_0.2_instruct_7b_en/2",
     },
 }

keras_hub/src/models/mistral/mistral_transformer_decoder.py

@@ -215,7 +215,8 @@ class MistralTransformerDecoder(keras.layers.Layer):
         # Mistral uses a banded attention mask if sliding window is not None
         if self.sliding_window is not None:
             # Below is a workaround for `ops.triu` for Keras 2.
-            # TODO(tirthasheshpatel): Use `ops.triu` once Keras 2 support is removed.
+            # TODO(tirthasheshpatel): Use `ops.triu` once Keras 2 support is
+            # removed.
             # causal_mask = ops.triu(causal_mask, k=-self.sliding_window)
             i = ops.arange(output_length)[:, None] + cache_update_index
             j = ops.arange(input_length)[None, :]

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)

keras_hub/src/models/{mix_transformer/mix_transformer_backbone.py → mit/mit_backbone.py}

@@ -1,28 +1,35 @@
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
 import keras
 import numpy as np
 from keras import ops
 
 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")
 class MiTBackbone(FeaturePyramidBackbone):
     def __init__(
         self,
-        depths,
+        layerwise_depths,
         num_layers,
-        blockwise_num_heads,
-        blockwise_sr_ratios,
+        layerwise_num_heads,
+        layerwise_sr_ratios,
         max_drop_path_rate,
-        patch_sizes,
-        strides,
+        layerwise_patch_sizes,
+        layerwise_strides,
         image_shape=(None, None, 3),
         hidden_dims=None,
         **kwargs,
@@ -36,12 +43,12 @@ class MiTBackbone(FeaturePyramidBackbone):
             https://github.com/DavidLandup0/deepvision/tree/main/deepvision/models/classification/mix_transformer)
 
         Args:
-            depths: The number of transformer encoders to be used per layer in the
-                network.
+            layerwise_depths: The number of transformer encoders to be used per
+                layer in the network.
             num_layers: int. The number of Transformer layers.
-            blockwise_num_heads: list of integers, the number of heads to use
+            layerwise_num_heads: list of integers, the number of heads to use
             in the attention computation for each layer.
-            blockwise_sr_ratios: list of integers, the sequence reduction
+            layerwise_sr_ratios: list of integers, the sequence reduction
                 ratio to perform for each layer on the sequence before key and
                 value projections. If set to > 1, a `Conv2D` layer is used to
                 reduce the length of the sequence.
@@ -51,7 +58,8 @@ class MiTBackbone(FeaturePyramidBackbone):
             image_shape: optional shape tuple, defaults to (None, None, 3).
             hidden_dims: the embedding dims per hierarchical layer, used as
                 the levels of the feature pyramid.
-            patch_sizes: list of integers, the patch_size to apply for each layer.
+            patch_sizes: list of integers, the patch_size to apply for each
+                layer.
             strides: list of integers, stride to apply for each layer.
 
         Examples:
@@ -61,7 +69,7 @@ class MiTBackbone(FeaturePyramidBackbone):
         ```python
         images = np.ones(shape=(1, 96, 96, 3))
         labels = np.zeros(shape=(1, 96, 96, 1))
-        backbone = keras_hub.models.MiTBackbone.from_preset("mit_b0_imagenet")
+        backbone = keras_hub.models.MiTBackbone.from_preset("mit_b0_ade20k_512")
 
         # Evaluate model
         model(images)
@@ -75,7 +83,10 @@ class MiTBackbone(FeaturePyramidBackbone):
         model.fit(images, labels, epochs=3)
         ```
         """
-        dpr = [x for x in np.linspace(0.0, max_drop_path_rate, sum(depths))]
+        dpr = [
+            x
+            for x in np.linspace(0.0, max_drop_path_rate, sum(layerwise_depths))
+        ]
 
         # === Layers ===
         cur = 0
@@ -86,8 +97,8 @@ class MiTBackbone(FeaturePyramidBackbone):
         for i in range(num_layers):
             patch_embed_layer = OverlappingPatchingAndEmbedding(
                 project_dim=hidden_dims[i],
-                patch_size=patch_sizes[i],
-                stride=strides[i],
+                patch_size=layerwise_patch_sizes[i],
+                stride=layerwise_strides[i],
                 name=f"patch_and_embed_{i}",
             )
             patch_embedding_layers.append(patch_embed_layer)
@@ -95,16 +106,16 @@ class MiTBackbone(FeaturePyramidBackbone):
             transformer_block = [
                 HierarchicalTransformerEncoder(
                     project_dim=hidden_dims[i],
-                    num_heads=blockwise_num_heads[i],
-                    sr_ratio=blockwise_sr_ratios[i],
+                    num_heads=layerwise_num_heads[i],
+                    sr_ratio=layerwise_sr_ratios[i],
                     drop_prob=dpr[cur + k],
                     name=f"hierarchical_encoder_{i}_{k}",
                 )
-                for k in range(depths[i])
+                for k in range(layerwise_depths[i])
             ]
             transformer_blocks.append(transformer_block)
-            cur += depths[i]
-            layer_norms.append(keras.layers.LayerNormalization())
+            cur += layerwise_depths[i]
+            layer_norms.append(keras.layers.LayerNormalization(epsilon=1e-5))
 
         # === Functional Model ===
         image_input = keras.layers.Input(shape=image_shape)
@@ -113,7 +124,7 @@ class MiTBackbone(FeaturePyramidBackbone):
         for i in range(num_layers):
             # Compute new height/width after the `proj`
             # call in `OverlappingPatchingAndEmbedding`
-            stride = strides[i]
+            stride = layerwise_strides[i]
             new_height, new_width = (
                 int(ops.shape(x)[1] / stride),
                 int(ops.shape(x)[2] / stride),
@@ -131,30 +142,30 @@ class MiTBackbone(FeaturePyramidBackbone):
         super().__init__(inputs=image_input, outputs=x, **kwargs)
 
         # === Config ===
-        self.depths = depths
+        self.layerwise_depths = layerwise_depths
         self.image_shape = image_shape
         self.hidden_dims = hidden_dims
         self.pyramid_outputs = pyramid_outputs
         self.num_layers = num_layers
-        self.blockwise_num_heads = blockwise_num_heads
-        self.blockwise_sr_ratios = blockwise_sr_ratios
+        self.layerwise_num_heads = layerwise_num_heads
+        self.layerwise_sr_ratios = layerwise_sr_ratios
         self.max_drop_path_rate = max_drop_path_rate
-        self.patch_sizes = patch_sizes
-        self.strides = strides
+        self.layerwise_patch_sizes = layerwise_patch_sizes
+        self.layerwise_strides = layerwise_strides
 
     def get_config(self):
         config = super().get_config()
         config.update(
             {
-                "depths": self.depths,
+                "layerwise_depths": self.layerwise_depths,
                 "hidden_dims": self.hidden_dims,
                 "image_shape": self.image_shape,
                 "num_layers": self.num_layers,
-                "blockwise_num_heads": self.blockwise_num_heads,
-                "blockwise_sr_ratios": self.blockwise_sr_ratios,
+                "layerwise_num_heads": self.layerwise_num_heads,
+                "layerwise_sr_ratios": self.layerwise_sr_ratios,
                 "max_drop_path_rate": self.max_drop_path_rate,
-                "patch_sizes": self.patch_sizes,
-                "strides": self.strides,
+                "layerwise_patch_sizes": self.layerwise_patch_sizes,
+                "layerwise_strides": self.layerwise_strides,
             }
         )
         return config

keras_hub/src/models/mit/mit_image_classifier.py

@@ -0,0 +1,12 @@
+from keras_hub.src.api_export import keras_hub_export
+from keras_hub.src.models.image_classifier import ImageClassifier
+from keras_hub.src.models.mit.mit_backbone import MiTBackbone
+from keras_hub.src.models.mit.mit_image_classifier_preprocessor import (
+    MiTImageClassifierPreprocessor,
+)
+
+
+@keras_hub_export("keras_hub.models.MiTImageClassifier")
+class MiTImageClassifier(ImageClassifier):
+    backbone_cls = MiTBackbone
+    preprocessor_cls = MiTImageClassifierPreprocessor

keras_hub/src/models/mit/mit_image_classifier_preprocessor.py

@@ -0,0 +1,12 @@
+from keras_hub.src.api_export import keras_hub_export
+from keras_hub.src.models.image_classifier_preprocessor import (
+    ImageClassifierPreprocessor,
+)
+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")
+class MiTImageClassifierPreprocessor(ImageClassifierPreprocessor):
+    backbone_cls = MiTBackbone
+    image_converter_cls = MiTImageConverter

keras_hub/src/models/mit/mit_image_converter.py

@@ -0,0 +1,8 @@
+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.mit import MiTBackbone
+
+
+@keras_hub_export("keras_hub.layers.MiTImageConverter")
+class MiTImageConverter(ImageConverter):
+    backbone_cls = MiTBackbone

keras_hub/src/models/{mix_transformer/mix_transformer_layers.py → mit/mit_layers.py}

@@ -28,19 +28,23 @@ class OverlappingPatchingAndEmbedding(keras.layers.Layer):
         self.patch_size = patch_size
         self.stride = stride
 
+        padding_size = self.patch_size // 2
+
+        self.padding = keras.layers.ZeroPadding2D(
+            padding=(padding_size, padding_size)
+        )
         self.proj = keras.layers.Conv2D(
             filters=project_dim,
             kernel_size=patch_size,
             strides=stride,
-            padding="same",
+            padding="valid",
         )
-        self.norm = keras.layers.LayerNormalization()
+        self.norm = keras.layers.LayerNormalization(epsilon=1e-5)
 
     def call(self, x):
+        x = self.padding(x)
         x = self.proj(x)
-        # B, H, W, C
-        shape = x.shape
-        x = ops.reshape(x, (-1, shape[1] * shape[2], shape[3]))
+        x = ops.reshape(x, (-1, x.shape[1] * x.shape[2], x.shape[3]))
         x = self.norm(x)
         return x
 
@@ -76,7 +80,8 @@ class HierarchicalTransformerEncoder(keras.layers.Layer):
             `LayerNormalization` layers. Defaults to `1e-06`
         sr_ratio: integer, the ratio to use within
             `SegFormerMultiheadAttention`. If set to > 1, a `Conv2D`
-             layer is used to reduce the length of the sequence. Defaults to `1`.
+            layer is used to reduce the length of the sequence.
+            Defaults to `1`.
     """
 
     def __init__(
@@ -179,20 +184,21 @@ class SegFormerMultiheadAttention(keras.layers.Layer):
         self.k = keras.layers.Dense(project_dim)
         self.v = keras.layers.Dense(project_dim)
         self.proj = keras.layers.Dense(project_dim)
+        self.dropout = keras.layers.Dropout(0.1)
+        self.proj_drop = keras.layers.Dropout(0.1)
 
         if sr_ratio > 1:
             self.sr = keras.layers.Conv2D(
                 filters=project_dim,
                 kernel_size=sr_ratio,
                 strides=sr_ratio,
-                padding="same",
             )
-            self.norm = keras.layers.LayerNormalization()
+            self.norm = keras.layers.LayerNormalization(epsilon=1e-5)
 
     def call(self, x):
         input_shape = ops.shape(x)
         H, W = int(math.sqrt(input_shape[1])), int(math.sqrt(input_shape[1]))
-        B, C = input_shape[0], input_shape[2]
+        B, N, C = input_shape[0], input_shape[1], input_shape[2]
 
         q = self.q(x)
         q = ops.reshape(
@@ -208,12 +214,11 @@ class SegFormerMultiheadAttention(keras.layers.Layer):
 
         if self.sr_ratio > 1:
             x = ops.reshape(
-                ops.transpose(x, [0, 2, 1]),
+                x,
                 (B, H, W, C),
             )
             x = self.sr(x)
-            x = ops.reshape(x, [input_shape[0], input_shape[2], -1])
-            x = ops.transpose(x, [0, 2, 1])
+            x = ops.reshape(x, [B, -1, C])
             x = self.norm(x)
 
         k = self.k(x)
@@ -237,14 +242,16 @@ class SegFormerMultiheadAttention(keras.layers.Layer):
 
         attn = (q @ ops.transpose(k, [0, 1, 3, 2])) * self.scale
         attn = ops.nn.softmax(attn, axis=-1)
+        attn = self.dropout(attn)
 
         attn = attn @ v
         attn = ops.reshape(
             ops.transpose(attn, [0, 2, 1, 3]),
-            [input_shape[0], input_shape[1], input_shape[2]],
+            [B, N, C],
         )
 
         x = self.proj(attn)
+        x = self.proj_drop(x)
         return x