keras-hub-nightly 0.15.0.dev20240823171555__py3-none-any.whl → 0.16.0.dev2024092017__py3-none-any.whl

keras_hub/src/models/efficientnet/fusedmbconv.py

@@ -0,0 +1,229 @@
+# Copyright 2024 The KerasHub Authors
+#
+# 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
+
+BN_AXIS = 3
+
+CONV_KERNEL_INITIALIZER = {
+    "class_name": "VarianceScaling",
+    "config": {
+        "scale": 2.0,
+        "mode": "fan_out",
+        "distribution": "truncated_normal",
+    },
+}
+
+
+class FusedMBConvBlock(keras.layers.Layer):
+    """Implementation of the FusedMBConv block
+
+    Also known as a Fused Mobile Inverted Residual Bottleneck block from:
+        [EfficientNet-EdgeTPU: Creating Accelerator-Optimized Neural Networks with AutoML]
+        (https://ai.googleblog.com/2019/08/efficientnet-edgetpu-creating.html)
+        [EfficientNetV2: Smaller Models and Faster Training]
+        (https://arxiv.org/abs/2104.00298v3).
+
+    FusedMBConv blocks are based on MBConv blocks, and replace the depthwise and
+    1x1 output convolution blocks with a single 3x3 convolution block, fusing
+    them together - hence the name "FusedMBConv". Alongside MBConv blocks, they
+    can be used in mobile-oriented and efficient architectures, and are present
+    in architectures EfficientNet.
+
+    FusedMBConv blocks follow a narrow-wide-narrow structure - expanding a 1x1
+    convolution, performing Squeeze-Excitation and then applying a 3x3
+    convolution, which is a more efficient operation than conventional
+    wide-narrow-wide structures.
+
+    As they're frequently used for models to be deployed to edge devices,
+    they're implemented as a layer for ease of use and re-use.
+
+    Args:
+        input_filters: int, the number of input filters
+        output_filters: int, the number of output filters
+        expand_ratio: default 1, the ratio by which input_filters are multiplied
+            to expand the structure in the middle expansion phase
+        kernel_size: default 3, the kernel_size to apply to the expansion phase
+            convolutions
+        strides: default 1, the strides to apply to the expansion phase
+            convolutions
+        se_ratio: default 0.0, The filters used in the Squeeze-Excitation phase,
+            and are chosen as the maximum between 1 and input_filters*se_ratio
+        batch_norm_momentum: default 0.9, the BatchNormalization momentum
+        activation: default "swish", the activation function used between
+            convolution operations
+        dropout: float, the optional dropout rate to apply before the output
+            convolution, defaults to 0.2
+
+    Returns:
+        A tensor representing a feature map, passed through the FusedMBConv
+        block
+
+    Note:
+        Not intended to be used outside of the EfficientNet architecture.
+    """
+
+    def __init__(
+        self,
+        input_filters,
+        output_filters,
+        expand_ratio=1,
+        kernel_size=3,
+        strides=1,
+        se_ratio=0.0,
+        batch_norm_momentum=0.9,
+        activation="swish",
+        dropout=0.2,
+        **kwargs
+    ):
+        super().__init__(**kwargs)
+        self.input_filters = input_filters
+        self.output_filters = output_filters
+        self.expand_ratio = expand_ratio
+        self.kernel_size = kernel_size
+        self.strides = strides
+        self.se_ratio = se_ratio
+        self.batch_norm_momentum = batch_norm_momentum
+        self.activation = activation
+        self.dropout = dropout
+        self.filters = self.input_filters * self.expand_ratio
+        self.filters_se = max(1, int(input_filters * se_ratio))
+
+        self.conv1 = keras.layers.Conv2D(
+            filters=self.filters,
+            kernel_size=kernel_size,
+            strides=strides,
+            kernel_initializer=CONV_KERNEL_INITIALIZER,
+            padding="same",
+            data_format="channels_last",
+            use_bias=False,
+            name=self.name + "expand_conv",
+        )
+        self.bn1 = keras.layers.BatchNormalization(
+            axis=BN_AXIS,
+            momentum=self.batch_norm_momentum,
+            name=self.name + "expand_bn",
+        )
+        self.act = keras.layers.Activation(
+            self.activation, name=self.name + "expand_activation"
+        )
+
+        self.bn2 = keras.layers.BatchNormalization(
+            axis=BN_AXIS,
+            momentum=self.batch_norm_momentum,
+            name=self.name + "bn",
+        )
+
+        self.se_conv1 = keras.layers.Conv2D(
+            self.filters_se,
+            1,
+            padding="same",
+            activation=self.activation,
+            kernel_initializer=CONV_KERNEL_INITIALIZER,
+            name=self.name + "se_reduce",
+        )
+
+        self.se_conv2 = keras.layers.Conv2D(
+            self.filters,
+            1,
+            padding="same",
+            activation="sigmoid",
+            kernel_initializer=CONV_KERNEL_INITIALIZER,
+            name=self.name + "se_expand",
+        )
+
+        self.output_conv = keras.layers.Conv2D(
+            filters=self.output_filters,
+            kernel_size=1 if expand_ratio != 1 else kernel_size,
+            strides=1,
+            kernel_initializer=CONV_KERNEL_INITIALIZER,
+            padding="same",
+            data_format="channels_last",
+            use_bias=False,
+            name=self.name + "project_conv",
+        )
+
+        self.bn3 = keras.layers.BatchNormalization(
+            axis=BN_AXIS,
+            momentum=self.batch_norm_momentum,
+            name=self.name + "project_bn",
+        )
+
+        if self.dropout:
+            self.dropout_layer = keras.layers.Dropout(
+                self.dropout,
+                noise_shape=(None, 1, 1, 1),
+                name=self.name + "drop",
+            )
+
+    def build(self, input_shape):
+        if self.name is None:
+            self.name = keras.backend.get_uid("block0")
+
+    def call(self, inputs):
+        # Expansion phase
+        if self.expand_ratio != 1:
+            x = self.conv1(inputs)
+            x = self.bn1(x)
+            x = self.act(x)
+        else:
+            x = inputs
+
+        # Squeeze and excite
+        if 0 < self.se_ratio <= 1:
+            se = keras.layers.GlobalAveragePooling2D(
+                name=self.name + "se_squeeze"
+            )(x)
+            if BN_AXIS == 1:
+                se_shape = (self.filters, 1, 1)
+            else:
+                se_shape = (1, 1, self.filters)
+
+            se = keras.layers.Reshape(se_shape, name=self.name + "se_reshape")(
+                se
+            )
+
+            se = self.se_conv1(se)
+            se = self.se_conv2(se)
+
+            x = keras.layers.multiply([x, se], name=self.name + "se_excite")
+
+        # Output phase:
+        x = self.output_conv(x)
+        x = self.bn3(x)
+        if self.expand_ratio == 1:
+            x = self.act(x)
+
+        # Residual:
+        if self.strides == 1 and self.input_filters == self.output_filters:
+            if self.dropout:
+                x = self.dropout_layer(x)
+            x = keras.layers.Add(name=self.name + "add")([x, inputs])
+        return x
+
+    def get_config(self):
+        config = {
+            "input_filters": self.input_filters,
+            "output_filters": self.output_filters,
+            "expand_ratio": self.expand_ratio,
+            "kernel_size": self.kernel_size,
+            "strides": self.strides,
+            "se_ratio": self.se_ratio,
+            "batch_norm_momentum": self.batch_norm_momentum,
+            "activation": self.activation,
+            "dropout": self.dropout,
+        }
+
+        base_config = super().get_config()
+        return dict(list(base_config.items()) + list(config.items()))

keras_hub/src/models/efficientnet/mbconv.py

@@ -0,0 +1,238 @@
+# Copyright 2024 The KerasHub Authors
+#
+# 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
+
+BN_AXIS = 3
+
+CONV_KERNEL_INITIALIZER = {
+    "class_name": "VarianceScaling",
+    "config": {
+        "scale": 2.0,
+        "mode": "fan_out",
+        "distribution": "truncated_normal",
+    },
+}
+
+
+class MBConvBlock(keras.layers.Layer):
+    def __init__(
+        self,
+        input_filters,
+        output_filters,
+        expand_ratio=1,
+        kernel_size=3,
+        strides=1,
+        se_ratio=0.0,
+        batch_norm_momentum=0.9,
+        activation="swish",
+        dropout=0.2,
+        **kwargs
+    ):
+        """Implementation of the MBConv block
+
+        Also known as a Mobile Inverted Residual Bottleneck block from:
+            [MobileNetV2: Inverted Residuals and Linear Bottlenecks]
+            (https://arxiv.org/abs/1801.04381v4).
+
+        MBConv blocks are common blocks used in mobile-oriented and efficient
+        architectures, present in architectures such as MobileNet, EfficientNet,
+        MaxViT, etc.
+
+        MBConv blocks follow a narrow-wide-narrow structure - expanding a 1x1
+        convolution, applying depthwise convolution, and narrowing back to a 1x1
+        convolution, which is a more efficient operation than conventional
+        wide-narrow-wide structures.
+
+        As they're frequently used for models to be deployed to edge devices,
+        they're implemented as a layer for ease of use and re-use.
+
+        Args:
+            input_filters: int, the number of input filters
+            output_filters: int, the optional number of output filters after
+                Squeeze-Excitation
+            expand_ratio: default 1, the ratio by which input_filters are
+                multiplied to expand the structure in the middle expansion phase
+            kernel_size: default 3, the kernel_size to apply to the expansion
+                phase convolutions
+            strides: default 1, the strides to apply to the expansion phase
+                convolutions
+            se_ratio: default 0.0, Squeeze-Excitation happens before depthwise
+                convolution and before output convolution only if the se_ratio
+                is above 0. The filters used in this phase are chosen as the
+                maximum between 1 and input_filters*se_ratio
+            batch_norm_momentum: default 0.9, the BatchNormalization momentum
+            activation: default "swish", the activation function used between
+                convolution operations
+            dropout: float, the optional dropout rate to apply before the output
+                convolution, defaults to 0.2
+
+        Returns:
+            A tensor representing a feature map, passed through the MBConv
+            block
+
+
+        Note:
+            Not intended to be used outside of the EfficientNet architecture.
+        """
+
+        super().__init__(**kwargs)
+        self.input_filters = input_filters
+        self.output_filters = output_filters
+        self.expand_ratio = expand_ratio
+        self.kernel_size = kernel_size
+        self.strides = strides
+        self.se_ratio = se_ratio
+        self.batch_norm_momentum = batch_norm_momentum
+        self.activation = activation
+        self.dropout = dropout
+        self.filters = self.input_filters * self.expand_ratio
+        self.filters_se = max(1, int(input_filters * se_ratio))
+
+        self.conv1 = keras.layers.Conv2D(
+            filters=self.filters,
+            kernel_size=1,
+            strides=1,
+            kernel_initializer=CONV_KERNEL_INITIALIZER,
+            padding="same",
+            data_format="channels_last",
+            use_bias=False,
+            name=self.name + "expand_conv",
+        )
+        self.bn1 = keras.layers.BatchNormalization(
+            axis=BN_AXIS,
+            momentum=self.batch_norm_momentum,
+            name=self.name + "expand_bn",
+        )
+        self.act = keras.layers.Activation(
+            self.activation, name=self.name + "activation"
+        )
+        self.depthwise = keras.layers.DepthwiseConv2D(
+            kernel_size=self.kernel_size,
+            strides=self.strides,
+            depthwise_initializer=CONV_KERNEL_INITIALIZER,
+            padding="same",
+            data_format="channels_last",
+            use_bias=False,
+            name=self.name + "dwconv2",
+        )
+
+        self.bn2 = keras.layers.BatchNormalization(
+            axis=BN_AXIS,
+            momentum=self.batch_norm_momentum,
+            name=self.name + "bn",
+        )
+
+        self.se_conv1 = keras.layers.Conv2D(
+            self.filters_se,
+            1,
+            padding="same",
+            activation=self.activation,
+            kernel_initializer=CONV_KERNEL_INITIALIZER,
+            name=self.name + "se_reduce",
+        )
+
+        self.se_conv2 = keras.layers.Conv2D(
+            self.filters,
+            1,
+            padding="same",
+            activation="sigmoid",
+            kernel_initializer=CONV_KERNEL_INITIALIZER,
+            name=self.name + "se_expand",
+        )
+
+        self.output_conv = keras.layers.Conv2D(
+            filters=self.output_filters,
+            kernel_size=1 if expand_ratio != 1 else kernel_size,
+            strides=1,
+            kernel_initializer=CONV_KERNEL_INITIALIZER,
+            padding="same",
+            data_format="channels_last",
+            use_bias=False,
+            name=self.name + "project_conv",
+        )
+
+        self.bn3 = keras.layers.BatchNormalization(
+            axis=BN_AXIS,
+            momentum=self.batch_norm_momentum,
+            name=self.name + "project_bn",
+        )
+
+        if self.dropout:
+            self.dropout_layer = keras.layers.Dropout(
+                self.dropout,
+                noise_shape=(None, 1, 1, 1),
+                name=self.name + "drop",
+            )
+
+    def build(self, input_shape):
+        if self.name is None:
+            self.name = keras.backend.get_uid("block0")
+
+    def call(self, inputs):
+        # Expansion phase
+        if self.expand_ratio != 1:
+            x = self.conv1(inputs)
+            x = self.bn1(x)
+            x = self.act(x)
+        else:
+            x = inputs
+
+        # Depthwise conv
+        x = self.depthwise(x)
+        x = self.bn2(x)
+        x = self.act(x)
+
+        # Squeeze and excite
+        if 0 < self.se_ratio <= 1:
+            se = keras.layers.GlobalAveragePooling2D(
+                name=self.name + "se_squeeze"
+            )(x)
+            if BN_AXIS == 1:
+                se_shape = (self.filters, 1, 1)
+            else:
+                se_shape = (1, 1, self.filters)
+            se = keras.layers.Reshape(se_shape, name=self.name + "se_reshape")(
+                se
+            )
+
+            se = self.se_conv1(se)
+            se = self.se_conv2(se)
+
+            x = keras.layers.multiply([x, se], name=self.name + "se_excite")
+
+        # Output phase
+        x = self.output_conv(x)
+        x = self.bn3(x)
+
+        if self.strides == 1 and self.input_filters == self.output_filters:
+            if self.dropout:
+                x = self.dropout_layer(x)
+            x = keras.layers.Add(name=self.name + "add")([x, inputs])
+        return x
+
+    def get_config(self):
+        config = {
+            "input_filters": self.input_filters,
+            "output_filters": self.output_filters,
+            "expand_ratio": self.expand_ratio,
+            "kernel_size": self.kernel_size,
+            "strides": self.strides,
+            "se_ratio": self.se_ratio,
+            "batch_norm_momentum": self.batch_norm_momentum,
+            "activation": self.activation,
+            "dropout": self.dropout,
+        }
+        base_config = super().get_config()
+        return dict(list(base_config.items()) + list(config.items()))

keras_hub/src/models/electra/__init__.py

@@ -14,7 +14,6 @@
 
 from keras_hub.src.models.electra.electra_backbone import ElectraBackbone
 from keras_hub.src.models.electra.electra_presets import backbone_presets
-from keras_hub.src.models.electra.electra_tokenizer import ElectraTokenizer
 from keras_hub.src.utils.preset_utils import register_presets
 
-register_presets(backbone_presets, (ElectraBackbone, ElectraTokenizer))
+register_presets(backbone_presets, ElectraBackbone)

keras_hub/src/models/electra/electra_tokenizer.py

@@ -13,10 +13,16 @@
 # limitations under the License.
 
 from keras_hub.src.api_export import keras_hub_export
+from keras_hub.src.models.electra.electra_backbone import ElectraBackbone
 from keras_hub.src.tokenizers.word_piece_tokenizer import WordPieceTokenizer
 
 
-@keras_hub_export("keras_hub.models.ElectraTokenizer")
+@keras_hub_export(
+    [
+        "keras_hub.tokenizers.ElectraTokenizer",
+        "keras_hub.models.ElectraTokenizer",
+    ]
+)
 class ElectraTokenizer(WordPieceTokenizer):
     """A ELECTRA tokenizer using WordPiece subword segmentation.
 
@@ -60,45 +66,24 @@ class ElectraTokenizer(WordPieceTokenizer):
     ```
     """
 
+    backbone_cls = ElectraBackbone
+
     def __init__(
         self,
         vocabulary,
         lowercase=False,
-        special_tokens_in_strings=False,
         **kwargs,
     ):
-        self.cls_token = "[CLS]"
-        self.sep_token = "[SEP]"
-        self.pad_token = "[PAD]"
-        self.mask_token = "[MASK]"
+        self._add_special_token("[CLS]", "cls_token")
+        self._add_special_token("[SEP]", "sep_token")
+        self._add_special_token("[PAD]", "pad_token")
+        self._add_special_token("[MASK]", "mask_token")
+        # Also add `tokenizer.start_token` and `tokenizer.end_token` for
+        # compatibility with other tokenizers.
+        self._add_special_token("[CLS]", "start_token")
+        self._add_special_token("[SEP]", "end_token")
         super().__init__(
             vocabulary=vocabulary,
             lowercase=lowercase,
-            special_tokens=[
-                self.cls_token,
-                self.sep_token,
-                self.pad_token,
-                self.mask_token,
-            ],
-            special_tokens_in_strings=special_tokens_in_strings,
             **kwargs,
         )
-
-    def set_vocabulary(self, vocabulary):
-        super().set_vocabulary(vocabulary)
-
-        if vocabulary is not None:
-            self.cls_token_id = self.token_to_id(self.cls_token)
-            self.sep_token_id = self.token_to_id(self.sep_token)
-            self.pad_token_id = self.token_to_id(self.pad_token)
-            self.mask_token_id = self.token_to_id(self.mask_token)
-        else:
-            self.cls_token_id = None
-            self.sep_token_id = None
-            self.pad_token_id = None
-            self.mask_token_id = None
-
-    def get_config(self):
-        config = super().get_config()
-        del config["special_tokens"]  # Not configurable; set in __init__.
-        return config

keras_hub/src/models/f_net/__init__.py

@@ -14,7 +14,6 @@
 
 from keras_hub.src.models.f_net.f_net_backbone import FNetBackbone
 from keras_hub.src.models.f_net.f_net_presets import backbone_presets
-from keras_hub.src.models.f_net.f_net_tokenizer import FNetTokenizer
 from keras_hub.src.utils.preset_utils import register_presets
 
-register_presets(backbone_presets, (FNetBackbone, FNetTokenizer))
+register_presets(backbone_presets, FNetBackbone)

keras_hub/src/models/f_net/f_net_masked_lm_preprocessor.py

@@ -13,17 +13,16 @@
 # limitations under the License.
 
 import keras
-from absl import logging
 
 from keras_hub.src.api_export import keras_hub_export
-from keras_hub.src.layers.preprocessing.masked_lm_mask_generator import (
-    MaskedLMMaskGenerator,
-)
-from keras_hub.src.models.f_net.f_net_preprocessor import FNetPreprocessor
+from keras_hub.src.models.f_net.f_net_backbone import FNetBackbone
+from keras_hub.src.models.f_net.f_net_tokenizer import FNetTokenizer
+from keras_hub.src.models.masked_lm_preprocessor import MaskedLMPreprocessor
+from keras_hub.src.utils.tensor_utils import preprocessing_function
 
 
 @keras_hub_export("keras_hub.models.FNetMaskedLMPreprocessor")
-class FNetMaskedLMPreprocessor(FNetPreprocessor):
+class FNetMaskedLMPreprocessor(MaskedLMPreprocessor):
     """FNet preprocessing for the masked language modeling task.
 
     This preprocessing layer will prepare inputs for a masked language modeling
@@ -119,78 +118,13 @@ class FNetMaskedLMPreprocessor(FNetPreprocessor):
     ```
     """
 
-    def __init__(
-        self,
-        tokenizer,
-        sequence_length=512,
-        truncate="round_robin",
-        mask_selection_rate=0.15,
-        mask_selection_length=96,
-        mask_token_rate=0.8,
-        random_token_rate=0.1,
-        **kwargs,
-    ):
-        super().__init__(
-            tokenizer,
-            sequence_length=sequence_length,
-            truncate=truncate,
-            **kwargs,
-        )
-        self.mask_selection_rate = mask_selection_rate
-        self.mask_selection_length = mask_selection_length
-        self.mask_token_rate = mask_token_rate
-        self.random_token_rate = random_token_rate
-        self.masker = None
-
-    def build(self, input_shape):
-        super().build(input_shape)
-        # Defer masker creation to `build()` so that we can be sure tokenizer
-        # assets have loaded when restoring a saved model.
-        self.masker = MaskedLMMaskGenerator(
-            mask_selection_rate=self.mask_selection_rate,
-            mask_selection_length=self.mask_selection_length,
-            mask_token_rate=self.mask_token_rate,
-            random_token_rate=self.random_token_rate,
-            vocabulary_size=self.tokenizer.vocabulary_size(),
-            mask_token_id=self.tokenizer.mask_token_id,
-            unselectable_token_ids=[
-                self.tokenizer.cls_token_id,
-                self.tokenizer.sep_token_id,
-                self.tokenizer.pad_token_id,
-            ],
-        )
-
-    def get_config(self):
-        config = super().get_config()
-        config.update(
-            {
-                "mask_selection_rate": self.mask_selection_rate,
-                "mask_selection_length": self.mask_selection_length,
-                "mask_token_rate": self.mask_token_rate,
-                "random_token_rate": self.random_token_rate,
-            }
-        )
-        return config
+    backbone_cls = FNetBackbone
+    tokenizer_cls = FNetTokenizer
 
+    @preprocessing_function
     def call(self, x, y=None, sample_weight=None):
-        if y is not None or sample_weight is not None:
-            logging.warning(
-                f"{self.__class__.__name__} generates `y` and `sample_weight` "
-                "based on your input data, but your data already contains `y` "
-                "or `sample_weight`. Your `y` and `sample_weight` will be "
-                "ignored."
-            )
-        x = super().call(x)
-        token_ids, segment_ids = (
-            x["token_ids"],
-            x["segment_ids"],
-        )
-        masker_outputs = self.masker(token_ids)
-        x = {
-            "token_ids": masker_outputs["token_ids"],
-            "segment_ids": segment_ids,
-            "mask_positions": masker_outputs["mask_positions"],
-        }
-        y = masker_outputs["mask_ids"]
-        sample_weight = masker_outputs["mask_weights"]
+        output = super().call(x, y=y, sample_weight=sample_weight)
+        x, y, sample_weight = keras.utils.unpack_x_y_sample_weight(output)
+        # FNet has not padding mask.
+        del x["padding_mask"]
         return keras.utils.pack_x_y_sample_weight(x, y, sample_weight)