mttf 1.3.6__py3-none-any.whl

mt/keras_src/applications_src/mobilevit.py

@@ -0,0 +1,292 @@
+# 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>`_.
+"""
+
+
+import tensorflow as tf
+from mt import tp, tfc
+
+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_src/base.py

@@ -0,0 +1,28 @@
+"""Determines the working Keras 2 from the system to be used by mt.keras."""
+
+from packaging.version import Version
+import tensorflow as tf
+
+tf_ver = Version(tf.__version__)
+if tf_ver >= Version("2.16"):
+    try:
+        import tf_keras
+    except:
+        raise ImportError(
+            f"mt.keras can only work with Keras 2. You have TF version {tf_ver}. Please install tf_keras."
+        )
+    keras_version = tf_keras.__version__
+    keras_source = "tf_keras"
+else:
+    try:
+        import keras
+
+        kr_ver = Version(keras.__version__)
+    except ImportError:
+        kr_ver = None
+    if kr_ver is None or kr_ver >= Version("3.0"):
+        keras_version = tf.__version__
+        keras_source = "tensorflow.python"
+    else:
+        keras_version = keras.__version__
+        keras_source = "keras"

mt/keras_src/constraints_src/__init__.py

@@ -0,0 +1,14 @@
+from .. import constraints as _constraints
+
+for _x, _y in _constraints.__dict__.items():
+    if _x.startswith("_"):
+        continue
+    globals()[_x] = _y
+__doc__ = _constraints.__doc__
+
+from .center_around import *
+
+
+__api__ = [
+    "CenterAround",
+]

mt/keras_src/constraints_src/center_around.py

@@ -0,0 +1,19 @@
+from .. import constraints
+
+
+class CenterAround(constraints.Constraint):
+    """Constrains the last axis to have values centered around `ref_value`."""
+
+    def __init__(self, ref_value: float = 0.0):
+        self.ref_value = ref_value
+
+    def __call__(self, w):
+        import tensorflow as tf
+
+        mean = tf.reduce_mean(w, axis=-1, keepdims=True)
+        ref_mean = mean - self.ref_value
+        ref_mean = tf.expand_dims(ref_mean, -1)
+        return w - ref_mean
+
+    def get_config(self):
+        return {"ref_value": self.ref_value}

mt/keras_src/layers_src/__init__.py

@@ -0,0 +1,43 @@
+from .. import layers as _layers
+
+for _x, _y in _layers.__dict__.items():
+    if _x.startswith("_"):
+        continue
+    globals()[_x] = _y
+__doc__ = _layers.__doc__
+
+from .identical import *
+from .floor import *
+from .var_regularizer import *
+from .simple_mha import *
+from .image_sizing import *
+from .counter import Counter
+from .normed_conv2d import NormedConv2D
+from .soft_bend import SoftBend
+from .transformer_encoder import MTTransformerEncoder
+from .utils import *
+
+
+__api__ = [
+    "Identical",
+    "Floor",
+    "VarianceRegularizer",
+    "SimpleMHA2D",
+    "MHAPool2D",
+    "DUCLayer",
+    "Downsize2D",
+    "Upsize2D",
+    "Downsize2D_V2",
+    "Upsize2D_V2",
+    "Downsize2D_V3",
+    "Downsize2D_V4",
+    "DownsizeX2D",
+    "UpsizeX2D",
+    "DownsizeY2D",
+    "UpsizeY2D",
+    "Counter",
+    "conv2d",
+    "dense2d",
+    "SoftBend",
+    "MTTransformerEncoder",
+]

mt/keras_src/layers_src/counter.py

@@ -0,0 +1,27 @@
+import tensorflow as tf
+from .. import layers, initializers
+
+
+class Counter(layers.Layer):
+    """A layer that counts from 0 during training and does nothing during inference."""
+
+    def build(self, input_shape):
+        initializer = initializers.Constant(value=0.0)
+        self.counter = self.add_weight(
+            name="counter", shape=(1,), initializer=initializer
+        )
+        self.incrementor = tf.constant([1.0])
+
+    def call(self, x, training: bool = False):
+        if training:
+            self.counter.assign_add(self.incrementor)
+        y = tf.reshape(x, [-1])[:1]
+        y = tf.stop_gradient(y) * 0.0
+        return self.counter + y
+
+    call.__doc__ = layers.Layer.call.__doc__
+
+    def compute_output_shape(self, input_shape):
+        return (1,)
+
+    compute_output_shape.__doc__ = layers.Layer.compute_output_shape.__doc__

mt/keras_src/layers_src/floor.py

@@ -0,0 +1,24 @@
+import tensorflow as tf
+from .. import layers
+
+
+@tf.custom_gradient
+def floor(x):
+    def grad(upstream):  # identity
+        return upstream
+
+    return tf.math.floor(x), grad
+
+
+class Floor(layers.Layer):
+    """TensorFlow floor but gradient is identity."""
+
+    def call(self, x):
+        return floor(x)
+
+    call.__doc__ = layers.Layer.call.__doc__
+
+    def compute_output_shape(self, input_shape):
+        return input_shape
+
+    compute_output_shape.__doc__ = layers.Layer.compute_output_shape.__doc__

mt/keras_src/layers_src/identical.py

@@ -0,0 +1,15 @@
+from .. import layers
+
+
+class Identical(layers.Layer):
+    """An identical layer, mainly for renaming purposes."""
+
+    def call(self, x):
+        return x
+
+    call.__doc__ = layers.Layer.call.__doc__
+
+    def compute_output_shape(self, input_shape):
+        return input_shape
+
+    compute_output_shape.__doc__ = layers.Layer.compute_output_shape.__doc__