mttf 1.1.10__py3-none-any.whl → 1.1.12__py3-none-any.whl

mt/keras/applications/__init__.py

@@ -0,0 +1,29 @@
+from ..base import applications as _applications
+
+for _x, _y in _applications.__dict__.items():
+    if _x.startswith("_"):
+        continue
+    globals()[_x] = _y
+__doc__ = _applications.__doc__
+
+from .mobilenet_v3_split import (
+    MobileNetV3Input,
+    MobileNetV3Parser,
+    MobileNetV3Block,
+    MobileNetV3Mixer,
+    MobileNetV3Output,
+    MobileNetV3Split,
+)
+
+from .mobilevit import create_mobilevit
+
+
+__api__ = [
+    "MobileNetV3Input",
+    "MobileNetV3Parser",
+    "MobileNetV3Block",
+    "MobileNetV3Mixer",
+    "MobileNetV3Output",
+    "MobileNetV3Split",
+    "create_mobilevit",
+]

mt/keras/applications/mobilenet_v3_split.py

@@ -0,0 +1,540 @@
+# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+#
+# 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
+#
+#     http://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.
+# ==============================================================================
+# pylint: disable=invalid-name
+# pylint: disable=missing-function-docstring
+"""MobileNet v3 models split into 5 submodels.
+
+The MobileNetV3 model is split into 5 parts:
+
+  - The input parser block that downsamples once (:func:`MobileNetV3Parser`).
+  - Block 0 to 3 that downample once for each block (:func:`MobileNetV3LargeBlock`
+    or :func:`MobileNetV3SmallBlock`). As of 2023/05/15, there's a possibility to have block 4
+    for MobileNetV3Large.
+  - The mixer block that turns the downsampled grid into a (1,1,feat_dim) batch
+    (:func:`MobileNetV3Mixer`).
+  - Optionally the output block that may or may not contain the clasification head
+    (:func:`MobileNetV3Output`).
+
+Input arguments follow those of MobileNetV3. One can also use :func:`MobileNetV3Split` to create
+a model of submodels that is theoretically equivalent to the original MobileNetV3 model. However,
+no pre-trained weights exist.
+"""
+
+
+from mt import tp, tfc
+
+from .. import base
+
+
+if base.keras_source == "tf_keras":
+    from tf_keras.applications.mobilenet_v3 import (
+        relu,
+        hard_swish,
+        _depth,
+        _inverted_res_block,
+    )
+    from tf_keras import backend, models, layers
+    from tensorflow.python.keras.utils import data_utils, layer_utils
+elif base.keras_source == "keras":
+    from keras.applications.mobilenet_v3 import (
+        relu,
+        hard_swish,
+        _depth,
+        _inverted_res_block,
+    )
+    from keras import backend, models
+    try:
+        from keras.layers import VersionAwareLayers
+
+        layers = VersionAwareLayers()
+    except ImportError:
+        try:
+            from keras import layers
+    from keras.utils import data_utils, layer_utils
+elif base.keras_source == "tensorflow.keras":
+    from tensorflow.keras.src.applications.mobilenet_v3 import (
+        relu,
+        hard_swish,
+        _depth,
+        _inverted_res_block,
+    )
+    from tensorflow.python.keras import backend, models
+    try:
+        from tensorflow.keras.layers import VersionAwareLayers
+
+        layers = VersionAwareLayers()
+    except ImportError:
+        try:
+            from tensorflow.keras import layers
+    from tensorflow.python.keras.utils import data_utils, layer_utils
+else:
+    raise ImportError(f"Unknown value '{keras_source}' for variable 'keras_source'.")
+
+from tensorflow.python.platform import tf_logging as logging
+
+
+def MobileNetV3Input(
+    input_shape=None,
+):
+    """Prepares a MobileNetV3 input layer."""
+
+    # If input_shape is None and input_tensor is None using standard shape
+    if input_shape is None:
+        input_shape = (None, None, 3)
+
+    if backend.image_data_format() == "channels_last":
+        row_axis, col_axis = (0, 1)
+    else:
+        row_axis, col_axis = (1, 2)
+    rows = input_shape[row_axis]
+    cols = input_shape[col_axis]
+    if rows and cols and (rows < 32 or cols < 32):
+        raise tfc.ModelSyntaxError(
+            f"Input size must be at least 32x32; got `input_shape={input_shape}`"
+        )
+
+    img_input = layers.Input(shape=input_shape)
+    return img_input
+
+
+def MobileNetV3Parser(
+    img_input,
+    model_type: str = "Large",  # only 'Small' or 'Large' are accepted
+    minimalistic=False,
+):
+    """Prepares a MobileNetV3 parser block."""
+
+    channel_axis = 1 if backend.image_data_format() == "channels_first" else -1
+
+    if minimalistic:
+        activation = relu
+    else:
+        activation = hard_swish
+
+    x = img_input
+    x = layers.Rescaling(scale=1.0 / 127.5, offset=-1.0)(x)
+    x = layers.Conv2D(
+        16, kernel_size=3, strides=(2, 2), padding="same", use_bias=False, name="Conv"
+    )(x)
+    x = layers.BatchNormalization(
+        axis=channel_axis, epsilon=1e-3, momentum=0.999, name="Conv/BatchNorm"
+    )(x)
+    x = activation(x)
+
+    # Create model.
+    model = models.Model(img_input, x, name=f"MobileNetV3{model_type}Parser")
+
+    return model
+
+
+def MobileNetV3SmallBlock(
+    block_id: int,  # only 0 to 3 are accepted here
+    input_tensor,  # input tensor for the block
+    alpha=1.0,
+    minimalistic=False,
+):
+    """Prepares a MobileNetV3Small downsampling block."""
+
+    def depth(d):
+        return _depth(d * alpha)
+
+    if minimalistic:
+        kernel = 3
+        activation = relu
+        se_ratio = None
+    else:
+        kernel = 5
+        activation = hard_swish
+        se_ratio = 0.25
+
+    x = input_tensor
+    if block_id == 0:
+        x = _inverted_res_block(x, 1, depth(16), 3, 2, se_ratio, relu, 0)
+    elif block_id == 1:
+        x = _inverted_res_block(x, 72.0 / 16, depth(24), 3, 2, None, relu, 1)
+        x = _inverted_res_block(x, 88.0 / 24, depth(24), 3, 1, None, relu, 2)
+    elif block_id == 2:
+        x = _inverted_res_block(x, 4, depth(40), kernel, 2, se_ratio, activation, 3)
+        x = _inverted_res_block(x, 6, depth(40), kernel, 1, se_ratio, activation, 4)
+        x = _inverted_res_block(x, 6, depth(40), kernel, 1, se_ratio, activation, 5)
+        x = _inverted_res_block(x, 3, depth(48), kernel, 1, se_ratio, activation, 6)
+        x = _inverted_res_block(x, 3, depth(48), kernel, 1, se_ratio, activation, 7)
+    else:
+        x = _inverted_res_block(x, 6, depth(96), kernel, 2, se_ratio, activation, 8)
+        x = _inverted_res_block(x, 6, depth(96), kernel, 1, se_ratio, activation, 9)
+        x = _inverted_res_block(x, 6, depth(96), kernel, 1, se_ratio, activation, 10)
+
+    # Create model.
+    model = models.Model(input_tensor, x, name=f"MobileNetV3SmallBlock{block_id}")
+
+    return model
+
+
+def MobileNetV3LargeBlock(
+    block_id: int,  # only 0 to 4 are accepted here. 4 is only available as of 2023/05/15
+    input_tensor,  # input tensor for the block
+    alpha=1.0,
+    minimalistic=False,
+):
+    """Prepares a MobileNetV3Large downsampling block."""
+
+    def depth(d):
+        return _depth(d * alpha)
+
+    if minimalistic:
+        kernel = 3
+        activation = relu
+        se_ratio = None
+    else:
+        kernel = 5
+        activation = hard_swish
+        se_ratio = 0.25
+
+    x = input_tensor
+    if block_id == 0:
+        x = _inverted_res_block(x, 1, depth(16), 3, 1, None, relu, 0)
+        x = _inverted_res_block(x, 4, depth(24), 3, 2, None, relu, 1)
+        x = _inverted_res_block(x, 3, depth(24), 3, 1, None, relu, 2)
+    elif block_id == 1:
+        x = _inverted_res_block(x, 3, depth(40), kernel, 2, se_ratio, relu, 3)
+        x = _inverted_res_block(x, 3, depth(40), kernel, 1, se_ratio, relu, 4)
+        x = _inverted_res_block(x, 3, depth(40), kernel, 1, se_ratio, relu, 5)
+    elif block_id == 2:
+        x = _inverted_res_block(x, 6, depth(80), 3, 2, None, activation, 6)
+        x = _inverted_res_block(x, 2.5, depth(80), 3, 1, None, activation, 7)
+        x = _inverted_res_block(x, 2.3, depth(80), 3, 1, None, activation, 8)
+        x = _inverted_res_block(x, 2.3, depth(80), 3, 1, None, activation, 9)
+        x = _inverted_res_block(x, 6, depth(112), 3, 1, se_ratio, activation, 10)
+        x = _inverted_res_block(x, 6, depth(112), 3, 1, se_ratio, activation, 11)
+    elif block_id == 3:
+        x = _inverted_res_block(x, 6, depth(160), kernel, 2, se_ratio, activation, 12)
+        x = _inverted_res_block(x, 6, depth(160), kernel, 1, se_ratio, activation, 13)
+        x = _inverted_res_block(x, 6, depth(160), kernel, 1, se_ratio, activation, 14)
+    else:
+        x = _inverted_res_block(x, 6, depth(320), kernel, 2, se_ratio, activation, 15)
+        x = _inverted_res_block(x, 6, depth(320), kernel, 1, se_ratio, activation, 16)
+        x = _inverted_res_block(x, 6, depth(320), kernel, 1, se_ratio, activation, 17)
+
+    # Create model.
+    model = models.Model(input_tensor, x, name=f"MobileNetV3LargeBlock{block_id}")
+
+    return model
+
+
+def MobileNetV3Mixer(
+    input_tensor,
+    params: tfc.MobileNetV3MixerParams,
+    last_point_ch,
+    alpha=1.0,
+    model_type: str = "Large",  # only 'Small' or 'Large' are accepted
+    minimalistic=False,
+):
+    """Prepares a MobileNetV3 mixer block."""
+
+    x = input_tensor
+    channel_axis = 1 if backend.image_data_format() == "channels_first" else -1
+
+    if params.variant == "mobilenet":
+
+        if minimalistic:
+            kernel = 3
+            activation = relu
+            se_ratio = None
+        else:
+            kernel = 5
+            activation = hard_swish
+            se_ratio = 0.25
+
+        last_conv_ch = _depth(backend.int_shape(x)[channel_axis] * 6)
+
+        # if the width multiplier is greater than 1 we
+        # increase the number of output channels
+        if alpha > 1.0:
+            last_point_ch = _depth(last_point_ch * alpha)
+        x = layers.Conv2D(
+            last_conv_ch, kernel_size=1, padding="same", use_bias=False, name="Conv_1"
+        )(x)
+        x = layers.BatchNormalization(
+            axis=channel_axis, epsilon=1e-3, momentum=0.999, name="Conv_1/BatchNorm"
+        )(x)
+        x = activation(x)
+        x = layers.GlobalAveragePooling2D()(x)
+        if channel_axis == 1:
+            x = layers.Reshape((last_conv_ch, 1, 1))(x)
+        else:
+            x = layers.Reshape((1, 1, last_conv_ch))(x)
+        x = layers.Conv2D(
+            last_point_ch, kernel_size=1, padding="same", use_bias=True, name="Conv_2"
+        )(x)
+        x = activation(x)
+    elif params.variant == "maxpool":
+        x = layers.GlobalMaxPool2D(x)
+    elif params.variant == "mhapool":
+        if backend.image_data_format() == "channels_first":
+            raise tfc.ModelSyntaxError(
+                "Mixer variant 'mhapool' requires channels_last image data format."
+            )
+
+        mhapool_params = params.mhapool_cascade_params
+        if not isinstance(mhapool_params, tfc.MHAPool2DCascadeParams):
+            raise tfc.ModelSyntaxError(
+                "Parameter 'params.mhapool_cascade_params' is not of type "
+                "mt.tfc.MHAPool2DCascadeParams. Got: {}.".format(type(mhapool_params))
+            )
+
+        from ..keras_layers import MHAPool2D
+
+        n_heads = mhapool_params.n_heads
+        k = 0
+        outputs = []
+        while True:
+            h = x.shape[1]
+            w = x.shape[2]
+
+            if h <= 1 and w <= 1:
+                break
+
+            c = x.shape[3]
+            key_dim = (c + n_heads - 1) // n_heads
+            value_dim = int(key_dim * mhapool_params.expansion_factor)
+            k += 1
+            block_name = f"MHAPool2DCascade_block{k}"
+            if k > mhapool_params.max_num_pooling_layers:  # GlobalMaxPool2D
+                x = layers.GlobalMaxPooling2D(
+                    keepdims=True, name=block_name + "/GlobalMaxPool"
+                )(x)
+            else:  # MHAPool2D
+                x = layers.LayerNormalization()(x)
+                if h <= 2 and w <= 2:
+                    activation = mhapool_params.final_activation
+                else:
+                    activation = mhapool_params.activation
+                x = MHAPool2D(
+                    n_heads,
+                    key_dim,
+                    value_dim=value_dim,
+                    pooling=mhapool_params.pooling,
+                    dropout=mhapool_params.dropout,
+                    name=block_name + "/MHAPool",
+                )(x)
+
+            if mhapool_params.output_all:
+                outputs.append(x)
+            else:
+                outputs = [x]
+    else:
+        raise tfc.ModelSyntaxError(
+            "Unknown mixer variant: '{}'.".format(params.variant)
+        )
+
+    # Create model.
+    model = models.Model(
+        input_tensor, outputs, name="MobileNetV3{}Mixer".format(model_type)
+    )
+
+    return model
+
+
+def MobileNetV3Output(
+    input_tensor,
+    model_type: str = "Large",  # only 'Small' or 'Large' are accepted
+    include_top=True,
+    classes=1000,
+    pooling=None,
+    dropout_rate=0.2,
+    classifier_activation="softmax",
+):
+    """Prepares a MobileNetV3 output block."""
+
+    x = input_tensor
+    if include_top:
+        if dropout_rate > 0:
+            x = layers.Dropout(dropout_rate)(x)
+        x = layers.Conv2D(classes, kernel_size=1, padding="same", name="Logits")(x)
+        x = layers.Flatten()(x)
+        x = layers.Activation(activation=classifier_activation, name="Predictions")(x)
+    else:
+        if pooling == "avg":
+            x = layers.GlobalAveragePooling2D(name="avg_pool")(x)
+        elif pooling == "max":
+            x = layers.GlobalMaxPooling2D(name="max_pool")(x)
+        else:
+            return None
+
+    # Create model.
+    model = models.Model(input_tensor, x, name=f"MobileNetV3{model_type}Output")
+
+    return model
+
+
+def MobileNetV3Split(
+    input_shape=None,
+    alpha: float = 1.0,
+    model_type: str = "Large",
+    max_n_blocks: int = 6,
+    minimalistic: bool = False,
+    mixer_params: tp.Optional[tfc.MobileNetV3MixerParams] = None,
+    include_top: bool = True,
+    pooling=None,
+    classes: int = 1000,
+    dropout_rate: float = 0.2,
+    classifier_activation="softmax",
+    output_all: bool = False,
+    name: tp.Optional[str] = None,
+):
+    """Prepares a model of submodels which is equivalent to a MobileNetV3 model.
+
+    Parameters
+    ----------
+    input_shape : tuple
+        Optional shape tuple, to be specified if you would like to use a model with an input image
+        resolution that is not (224, 224, 3). It should have exactly 3 inputs channels
+        (224, 224, 3). You can also omit this option if you would like to infer input_shape from an
+        input_tensor. If you choose to include both input_tensor and input_shape then input_shape
+        will be used if they match, if the shapes do not match then we will throw an error. E.g.
+        `(160, 160, 3)` would be one valid value.
+    alpha : float
+        controls the width of the network. This is known as the depth multiplier in the MobileNetV3
+        paper, but the name is kept for consistency with MobileNetV1 in Keras.
+        - If `alpha` < 1.0, proportionally decreases the number
+            of filters in each layer.
+        - If `alpha` > 1.0, proportionally increases the number
+            of filters in each layer.
+        - If `alpha` = 1, default number of filters from the paper
+            are used at each layer.
+          the mobilenetv3 alpha value
+    model_type : {'Small', 'Large'}
+        whether it is the small variant or the large variant
+    max_n_blocks : int
+        the maximum number of blocks in the backbone. It is further constrained by the actual
+        maximum number of blocks that the variant can implement.
+    minimalistic : bool
+        In addition to large and small models this module also contains so-called minimalistic
+        models, these models have the same per-layer dimensions characteristic as MobilenetV3
+        however, they do not utilize any of the advanced blocks (squeeze-and-excite units,
+        hard-swish, and 5x5 convolutions). While these models are less efficient on CPU, they
+        are much more performant on GPU/DSP.
+    mixer_params : mt.tfc.MobileNetV3MixerParams, optional
+        parameters for defining the mixer block
+    include_top : bool, default True
+        whether to include the fully-connected layer at the top of the network. Only valid if
+        `mixer_params` is not null.
+    pooling : str, optional
+        Optional pooling mode for feature extraction when `include_top` is False and
+        `mixer_params` is not null.
+        - `None` means that the output of the model will be the 4D tensor output of the last
+          convolutional block.
+        - `avg` means that global average pooling will be applied to the output of the last
+          convolutional block, and thus the output of the model will be a 2D tensor.
+        - `max` means that global max pooling will be applied.
+    classes : int, optional
+        Optional number of classes to classify images into, only to be specified if `mixer_params`
+        is not null and `include_top` is True.
+    dropout_rate : float
+        fraction of the input units to drop on the last layer. Only to be specified if
+        `mixer_params` is not null and `include_top` is True.
+    classifier_activation : object
+        A `str` or callable. The activation function to use on the "top" layer. Ignored unless
+        `mixer_params` is not null and `include_top` is True. Set `classifier_activation=None` to
+        return the logits of the "top" layer. When loading pretrained weights,
+        `classifier_activation` can only be `None` or `"softmax"`.
+    output_all : bool
+        If True, the model returns the output tensor of every submodel other than the input layer.
+        Otherwise, it returns the output tensor of the last submodel.
+    name : str, optional
+        model name, if any. Default to 'MobileNetV3LargeSplit' or 'MobileNetV3SmallSplit'.
+
+    Returns
+    -------
+    tensorflow.keras.Model
+        the output MobileNetV3 model split into 5 submodels
+    """
+
+    input_layer = MobileNetV3Input(input_shape=input_shape)
+    input_block = MobileNetV3Parser(
+        input_layer,
+        model_type=model_type,
+        minimalistic=minimalistic,
+    )
+    x = input_block(input_layer)
+    outputs = [x]
+
+    num_blocks = 5 if model_type == "Large" else 4
+    if num_blocks > max_n_blocks:
+        num_blocks = max_n_blocks
+    for i in range(num_blocks):
+        if model_type == "Large":
+            block = MobileNetV3LargeBlock(i, x, alpha=alpha, minimalistic=minimalistic)
+        else:
+            block = MobileNetV3SmallBlock(i, x, alpha=alpha, minimalistic=minimalistic)
+        x = block(x)
+        if output_all:
+            outputs.append(x)
+        else:
+            outputs = [x]
+
+    if mixer_params is not None:
+        if not isinstance(mixer_params, tfc.MobileNetV3MixerParams):
+            raise tfc.ModelSyntaxError(
+                "Argument 'mixer_params' is not an instance of "
+                "mt.tfc.MobileNetV3MixerParams. Got: {}.".format(type(mixer_params))
+            )
+
+        if model_type == "Large":
+            last_point_ch = 1280
+        else:
+            last_point_ch = 1024
+        mixer_block = MobileNetV3Mixer(
+            x,
+            mixer_params,
+            last_point_ch,
+            alpha=alpha,
+            model_type=model_type,
+            minimalistic=minimalistic,
+        )
+        x = mixer_block(x)
+        if output_all:
+            if isinstance(x, (list, tuple)):
+                outputs.extend(x)
+            else:
+                outputs.append(x)
+        else:
+            if isinstance(x, (list, tuple)):
+                outputs = [x[-1]]
+            else:
+                outputs = [x]
+
+        output_block = MobileNetV3Output(
+            x,
+            model_type=model_type,
+            include_top=include_top,
+            classes=classes,
+            pooling=pooling,
+            dropout_rate=dropout_rate,
+            classifier_activation=classifier_activation,
+        )
+        if output_block is not None:
+            x = output_block(x)
+            if output_all:
+                outputs.append(x)
+            else:
+                outputs = [x]
+
+    # Create model.
+    if name is None:
+        name = f"MobilenetV3{model_type}Split"
+    model = models.Model(input_layer, outputs, name=name)
+
+    return model

mt/keras/applications/mobilevit.py

@@ -0,0 +1,291 @@
+# pylint: disable=invalid-name
+# pylint: disable=missing-function-docstring
+"""MobileViT model.
+
+Most of the code here has been ripped and updated off from the following
+`Keras tutorial <https://keras.io/examples/vision/mobilevit/>`_. Please refer
+to the `MobileViT ICLR2022 paper <https://arxiv.org/abs/2110.02178>`_ for more details.
+
+The paper authors' code is `here <https://github.com/apple/ml-cvnets>`_.
+"""
+
+
+from mt import tp, tfc, tf
+
+from .mobilenet_v3_split import (
+    MobileNetV3Input,
+    _inverted_res_block,
+    backend,
+    models,
+    layers,
+)
+
+
+def conv_block(x, filters=16, kernel_size=3, strides=2):
+    conv_layer = layers.Conv2D(
+        filters, kernel_size, strides=strides, activation=tf.nn.swish, padding="same"
+    )
+    return conv_layer(x)
+
+
+# Reference: https://git.io/JKgtC
+
+
+def inverted_residual_block(
+    x, expanded_channels, output_channels, strides=1, block_id=0
+):
+    if block_id == 0:
+        raise NotImplementedError(
+            "Zero block id for _inverted_res_block() is not implemented in MobileViT."
+        )
+
+    channel_axis = 1 if backend.image_data_format() == "channels_first" else -1
+    infilters = backend.int_shape(x)[channel_axis]
+
+    m = _inverted_res_block(
+        x,
+        expanded_channels // infilters,  # expansion
+        output_channels,  # filters
+        3,  # kernel_size
+        strides,  # stride
+        0,  # se_ratio
+        tf.nn.swish,  # activation
+        block_id,
+    )
+
+    return m
+
+
+# Reference:
+# https://keras.io/examples/vision/image_classification_with_vision_transformer/
+
+
+def mlp(x, hidden_units, dropout_rate):
+    for units in hidden_units:
+        x = layers.Dense(units, activation=tf.nn.swish)(x)
+        x = layers.Dropout(dropout_rate)(x)
+    return x
+
+
+def transformer_block(x, transformer_layers, projection_dim, num_heads=2):
+    for _ in range(transformer_layers):
+        # Layer normalization 1.
+        x1 = layers.LayerNormalization(epsilon=1e-6)(x)
+        # Create a multi-head attention layer.
+        attention_output = layers.MultiHeadAttention(
+            num_heads=num_heads, key_dim=projection_dim, dropout=0.1
+        )(x1, x1)
+        # Skip connection 1.
+        x2 = layers.Add()([attention_output, x])
+        # Layer normalization 2.
+        x3 = layers.LayerNormalization(epsilon=1e-6)(x2)
+        # MLP.
+        x3 = mlp(
+            x3,
+            hidden_units=[x.shape[-1] * 2, x.shape[-1]],
+            dropout_rate=0.1,
+        )
+        # Skip connection 2.
+        x = layers.Add()([x3, x2])
+
+    return x
+
+
+def mobilevit_block(x, num_blocks, projection_dim, strides=1):
+    cell_size = 2  # 2x2 for the Transformer block
+
+    # Local projection with convolutions.
+    local_features = conv_block(x, filters=projection_dim, strides=strides)
+    local_features = conv_block(
+        local_features, filters=projection_dim, kernel_size=1, strides=strides
+    )
+
+    if x.shape[1] % cell_size != 0:
+        raise tfc.ModelSyntaxError(
+            f"Input tensor must have height divisible by {cell_size}. Got {x.shape}."
+        )
+
+    if x.shape[2] % cell_size != 0:
+        raise tfc.ModelSyntaxError(
+            f"Input tensor must have width divisible by {cell_size}. Got {x.shape}."
+        )
+
+    # Unfold into patches and then pass through Transformers.
+    z = local_features  # (B,H,W,C)
+    z = layers.Reshape(
+        (
+            z.shape[1] // cell_size,
+            cell_size,
+            z.shape[2] // cell_size,
+            cell_size,
+            projection_dim,
+        )
+    )(
+        z
+    )  # (B,H/P,P,W/P,P,C)
+    z = tf.transpose(z, perm=[0, 2, 4, 1, 3, 5])  # (B,P,P,H/P,W/P,C)
+    non_overlapping_patches = layers.Reshape(
+        (cell_size * cell_size, z.shape[3] * z.shape[4], projection_dim)
+    )(
+        z
+    )  # (B,P*P,H*W/(P*P),C)
+    global_features = transformer_block(
+        non_overlapping_patches, num_blocks, projection_dim
+    )
+
+    # Fold into conv-like feature-maps.
+    z = layers.Reshape(
+        (
+            cell_size,
+            cell_size,
+            x.shape[1] // cell_size,
+            x.shape[2] // cell_size,
+            projection_dim,
+        )
+    )(
+        global_features
+    )  # (B,P,P,H/P,W/P,C)
+    z = tf.transpose(z, perm=[0, 3, 1, 4, 2, 5])  # (B,H/P,P,W/P,P,C)
+    folded_feature_map = layers.Reshape((x.shape[1], x.shape[2], projection_dim))(z)
+
+    # Apply point-wise conv -> concatenate with the input features.
+    folded_feature_map = conv_block(
+        folded_feature_map, filters=x.shape[-1], kernel_size=1, strides=strides
+    )
+    local_global_features = layers.Concatenate(axis=-1)([x, folded_feature_map])
+
+    # Fuse the local and global features using a convolution layer.
+    local_global_features = conv_block(
+        local_global_features, filters=projection_dim, strides=strides
+    )
+
+    return local_global_features
+
+
+def create_mobilevit(
+    input_shape=None,
+    model_type: str = "XXS",
+    output_all: bool = False,
+    name: tp.Optional[str] = None,
+):
+    """Prepares a model of submodels which is equivalent to a MobileNetV3 model.
+
+    Parameters
+    ----------
+    input_shape : tuple
+        Optional shape tuple, to be specified if you would like to use a model with an input image
+        resolution that is not (224, 224, 3). It should have exactly 3 inputs channels
+        (224, 224, 3). You can also omit this option if you would like to infer input_shape from an
+        input_tensor. If you choose to include both input_tensor and input_shape then input_shape
+        will be used if they match, if the shapes do not match then we will throw an error. E.g.
+        `(160, 160, 3)` would be one valid value.
+    model_type : {'XXS', 'XS', 'S'}
+        one of the 3 variants introduced in the paper
+    output_all : bool
+        If True, the model returns the output tensor of every block before down-sampling, other
+        than the input layer.  Otherwise, it returns the output tensor of the last block.
+    name : str, optional
+        model name, if any. Default to 'MobileViT<model_type>'.
+
+    Returns
+    -------
+    tensorflow.keras.Model
+        the output MobileViT model
+    """
+
+    model_type_id = ["XXS", "XS", "S"].index(model_type)
+
+    expansion_factor = 2 if model_type_id == 0 else 4
+
+    inputs = MobileNetV3Input(input_shape=input_shape)
+    x = layers.Rescaling(scale=1.0 / 255)(inputs)
+
+    # Initial conv-stem -> MV2 block.
+    x = conv_block(x, filters=16)
+    x = inverted_residual_block(
+        x,
+        expanded_channels=16 * expansion_factor,
+        output_channels=16 if model_type_id == 0 else 32,
+        block_id=1,
+    )
+    outputs = [x]
+
+    # Downsampling with MV2 block.
+    output_channels = [24, 48, 64][model_type_id]
+    x = inverted_residual_block(
+        x,
+        expanded_channels=16 * expansion_factor,
+        output_channels=output_channels,
+        strides=2,
+        block_id=2,
+    )
+    x = inverted_residual_block(
+        x,
+        expanded_channels=24 * expansion_factor,
+        output_channels=output_channels,
+        block_id=3,
+    )
+    x = inverted_residual_block(
+        x,
+        expanded_channels=24 * expansion_factor,
+        output_channels=output_channels,
+        block_id=4,
+    )
+    if output_all:
+        outputs.append(x)
+    else:
+        outputs = [x]
+
+    # First MV2 -> MobileViT block.
+    output_channels = [48, 64, 96][model_type_id]
+    projection_dim = [64, 96, 144][model_type_id]
+    x = inverted_residual_block(
+        x,
+        expanded_channels=48 * expansion_factor,
+        output_channels=output_channels,
+        strides=2,
+        block_id=5,
+    )
+    x = mobilevit_block(x, num_blocks=2, projection_dim=projection_dim)
+    if output_all:
+        outputs.append(x)
+    else:
+        outputs = [x]
+
+    # Second MV2 -> MobileViT block.
+    output_channels = [64, 80, 128][model_type_id]
+    projection_dim = [80, 120, 192][model_type_id]
+    x = inverted_residual_block(
+        x,
+        expanded_channels=64 * expansion_factor,
+        output_channels=output_channels,
+        strides=2,
+        block_id=6,
+    )
+    x = mobilevit_block(x, num_blocks=4, projection_dim=projection_dim)
+    if output_all:
+        outputs.append(x)
+    else:
+        outputs = [x]
+
+    # Third MV2 -> MobileViT block.
+    output_channels = [80, 96, 160][model_type_id]
+    projection_dim = [96, 144, 240][model_type_id]
+    x = inverted_residual_block(
+        x,
+        expanded_channels=80 * expansion_factor,
+        output_channels=output_channels,
+        strides=2,
+        block_id=7,
+    )
+    x = mobilevit_block(x, num_blocks=3, projection_dim=projection_dim)
+    filters = [320, 384, 640][model_type_id]
+    x = conv_block(x, filters=filters, kernel_size=1, strides=1)
+    if output_all:
+        outputs.append(x)
+    else:
+        outputs = [x]
+
+    if name is None:
+        name = f"MobileViT{model_type}"
+    return models.Model(inputs, outputs, name=name)

mt/keras/layers/__init__.py

@@ -4,7 +4,7 @@ for _x, _y in _layers.__dict__.items():
     if _x.startswith("_"):
         continue
     globals()[_x] = _y
-globals.__doc__ = _layers.__doc__
+__doc__ = _layers.__doc__
 
 from .identical import *
 from .floor import *

mt/tf/version.py

@@ -1,5 +1,5 @@
 MAJOR_VERSION = 1
 MINOR_VERSION = 1
-PATCH_VERSION = 10
+PATCH_VERSION = 12
 version = '{}.{}.{}'.format(MAJOR_VERSION, MINOR_VERSION, PATCH_VERSION)
 __all__  = ['MAJOR_VERSION', 'MINOR_VERSION', 'PATCH_VERSION', 'version']

{mttf-1.1.10.dist-info → mttf-1.1.12.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: mttf
-Version: 1.1.10
+Version: 1.1.12
 Summary: A package to detect and monkey-patch TensorFlow and Keras, for Minh-Tri Pham
 Home-page: https://github.com/inteplus/mttf
 Author: ['Minh-Tri Pham']

{mttf-1.1.10.dist-info → mttf-1.1.12.dist-info}/RECORD

@@ -1,7 +1,10 @@
 mt/keras/__init__.py,sha256=oM2xsSZTuGSUvFJSpntk8r7bzGcTytNvn-Apv_SBFh4,70
+mt/keras/applications/__init__.py,sha256=eTTWZQU9u0z7W0eDKo3GaVJvVCh2gDtjD8R7S-8ZWH0,593
+mt/keras/applications/mobilenet_v3_split.py,sha256=kpFEpE-93QYIt25ZbbGoXuhVzB-GyaCDJnXFCxd0-iU,19747
+mt/keras/applications/mobilevit.py,sha256=v-XzmFSwWAjN_nsp2ys0QINtoqcYD-qvNrSbKHkreQs,8914
 mt/keras/base/__init__.py,sha256=vIvuf3gjbpXC-rGpXLfIApLRUV6w1GUzVK8YJOgWLyk,327
 mt/keras/base/base.py,sha256=_B2sSUMlHOtGSAqQD1p5YD0raEDL4W0Bh3uKD6BXOJM,807
-mt/keras/layers/__init__.py,sha256=pqRphJpkmB1y59f17yPKEnOoy-_2tp5mPAD4fBY3pKQ,758
+mt/keras/layers/__init__.py,sha256=7TMhZig3bHtyDpiwnoXcwgAxw6e5xkAy664Z6qLQUrY,750
 mt/keras/layers/counter.py,sha256=J3__IXbaa7zp72a5P8FFi0bfftTHwa1xzzCwxCIU2gc,856
 mt/keras/layers/floor.py,sha256=4mSpmTrhM7VqTK85POkjC3OhaTZUNUF9knO7gTbSGtc,512
 mt/keras/layers/identical.py,sha256=AIqC36PxU9sXyF9rZuQ-5ObjnIjBiSIMHIb5MwqVdmY,361
@@ -14,7 +17,7 @@ mt/tf/__init__.py,sha256=M8xiJNdrAUJZgiZTOQOdfkehjO-CYzGpoxh5HVGBkms,338
 mt/tf/init.py,sha256=BUpw33uyA_DmeJjrN2fX4MIs8MynKxkwgc2oTGTqssU,1294
 mt/tf/mttf_version.py,sha256=ha53i-H9pE-crufFttUECgXHwPvam07zMKzApUts1Gs,206
 mt/tf/utils.py,sha256=Copl5VM0PpuFUchK-AcBuGO6QitDwHcEs4FruZb2GAI,2460
-mt/tf/version.py,sha256=xBDDlvCAF-wcYw_LU24Q5In5Tf5HJhO5hhybChnGc4A,207
+mt/tf/version.py,sha256=NG7gRVnIXelT8AS-NtYLCjdeQ9cag8SH1NfKOpmwVq0,207
 mt/tf/keras_applications/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 mt/tf/keras_applications/mobilenet_v3_split.py,sha256=1oPB3EX3k3c7iju9Ksuw9xyv32-mOKPs3uy2Mk5tLd8,19716
 mt/tf/keras_applications/mobilevit.py,sha256=VsKB_U-f9jwUEjpd0eq-YXL4rDsuAbKQ0yIzkbMfLzw,9949
@@ -23,14 +26,14 @@ mt/tfc/__init__.py,sha256=XFXAbmEi5uTAuZj797cBqw9ZBnVy6ptxP0TTzVauMNk,8446
 mt/tfg/__init__.py,sha256=6Ly2QImAyQTsg_ZszuAuK_L2n56v89Cix9yYmMVk0CM,304
 mt/tfp/__init__.py,sha256=AQkGCkmDRwswEt3qoOSpxe-fZekx78sHHBs2ZVz33gc,383
 mt/tfp/real_nvp.py,sha256=U9EmkXGqFcvtS2yeh5_RgbKlVKKlGFGklAb7Voyazz4,4440
-mttf-1.1.10.data/scripts/dmt_pipi.sh,sha256=NNsj4P332unHMqU4mAFjU9PQvxQ8TK5XQ42LC29IZY8,510
-mttf-1.1.10.data/scripts/dmt_twineu.sh,sha256=KZhcYwuCW0c36tWcOgCe7uxJmS08rz-J6YNY76Exy4M,193
-mttf-1.1.10.data/scripts/pipi.sh,sha256=kdo96bdaKq2QIa52Z4XFSiGPcbDm09SAU9cju6I2Lxo,289
-mttf-1.1.10.data/scripts/wml_nexus.py,sha256=kW0ju8_kdXc4jOjhdzKiMsFuO1MNpHmu87skrhu9SEg,1492
-mttf-1.1.10.data/scripts/wml_pipi.sh,sha256=CuidIcbuxyXSBNQqYRhCcSC8QbBaSGnQX0KAIFaIvKA,499
-mttf-1.1.10.data/scripts/wml_twineu.sh,sha256=av1JLN765oOWC5LPkv2eSWIVof26y60601tMGkuYdb8,180
-mttf-1.1.10.dist-info/licenses/LICENSE,sha256=e_JtcszdGZ2ZGfjcymTGrcxFj_9XPicZOVtnsrPvruk,1070
-mttf-1.1.10.dist-info/METADATA,sha256=TflB9A0LQvdZSg0o03NUnUjERgNXsQbuq0eM3zSF9tc,568
-mttf-1.1.10.dist-info/WHEEL,sha256=wXxTzcEDnjrTwFYjLPcsW_7_XihufBwmpiBeiXNBGEA,91
-mttf-1.1.10.dist-info/top_level.txt,sha256=WcqGFu9cV7iMZg09iam8eNxUvGpLSKKF2Iubf6SJVOo,3
-mttf-1.1.10.dist-info/RECORD,,
+mttf-1.1.12.data/scripts/dmt_pipi.sh,sha256=NNsj4P332unHMqU4mAFjU9PQvxQ8TK5XQ42LC29IZY8,510
+mttf-1.1.12.data/scripts/dmt_twineu.sh,sha256=KZhcYwuCW0c36tWcOgCe7uxJmS08rz-J6YNY76Exy4M,193
+mttf-1.1.12.data/scripts/pipi.sh,sha256=kdo96bdaKq2QIa52Z4XFSiGPcbDm09SAU9cju6I2Lxo,289
+mttf-1.1.12.data/scripts/wml_nexus.py,sha256=kW0ju8_kdXc4jOjhdzKiMsFuO1MNpHmu87skrhu9SEg,1492
+mttf-1.1.12.data/scripts/wml_pipi.sh,sha256=CuidIcbuxyXSBNQqYRhCcSC8QbBaSGnQX0KAIFaIvKA,499
+mttf-1.1.12.data/scripts/wml_twineu.sh,sha256=av1JLN765oOWC5LPkv2eSWIVof26y60601tMGkuYdb8,180
+mttf-1.1.12.dist-info/licenses/LICENSE,sha256=e_JtcszdGZ2ZGfjcymTGrcxFj_9XPicZOVtnsrPvruk,1070
+mttf-1.1.12.dist-info/METADATA,sha256=wqmEBPpwlrdEh_Asq2eXpphu-ZzFv2b4G1-drGywf84,568
+mttf-1.1.12.dist-info/WHEEL,sha256=wXxTzcEDnjrTwFYjLPcsW_7_XihufBwmpiBeiXNBGEA,91
+mttf-1.1.12.dist-info/top_level.txt,sha256=WcqGFu9cV7iMZg09iam8eNxUvGpLSKKF2Iubf6SJVOo,3
+mttf-1.1.12.dist-info/RECORD,,