keras-hub-nightly 0.23.0.dev202510090417__py3-none-any.whl → 0.23.0.dev202510110411__py3-none-any.whl

keras_hub/src/models/mobilenetv5/mobilenetv5_builder.py

@@ -0,0 +1,436 @@
+import re
+from copy import deepcopy
+
+from keras_hub.src.models.mobilenet.mobilenet_backbone import ConvBnActBlock
+from keras_hub.src.models.mobilenet.mobilenet_backbone import DepthwiseConvBlock
+from keras_hub.src.models.mobilenet.mobilenet_backbone import (
+    InvertedResidualBlock,
+)
+from keras_hub.src.models.mobilenet.util import adjust_channels
+from keras_hub.src.models.mobilenetv5.mobilenetv5_attention import (
+    MobileAttention,
+)
+from keras_hub.src.models.mobilenetv5.mobilenetv5_blocks import CondConvResidual
+from keras_hub.src.models.mobilenetv5.mobilenetv5_blocks import EdgeResidual
+from keras_hub.src.models.mobilenetv5.mobilenetv5_blocks import (
+    UniversalInvertedResidual,
+)
+from keras_hub.src.models.mobilenetv5.mobilenetv5_utils import parse_ksize
+from keras_hub.src.models.mobilenetv5.mobilenetv5_utils import round_channels
+
+
+def decode_block_str(block_str):
+    assert isinstance(block_str, str)
+    ops = block_str.split("_")
+    block_type = ops[0]
+    ops = ops[1:]
+    options = {}
+    skip = None
+    for op in ops:
+        if op == "noskip":
+            skip = False
+        elif op == "skip":
+            skip = True
+        elif op.startswith("n"):
+            key = op[0]
+            v = op[1:]
+            options[key] = v if v else "relu"
+        else:
+            splits = re.split(r"(\d.*)", op)
+            if len(splits) >= 2:
+                key, value = splits[:2]
+                options[key] = value
+
+    act_layer = options.get("n", "gelu")
+    num_repeat = int(options["r"])
+
+    block_args = dict(
+        block_type=block_type,
+        out_chs=int(options["c"]),
+        stride=int(options["s"]),
+        act_layer=act_layer,
+    )
+
+    if block_type == "ir":
+        block_args.update(
+            dict(
+                dw_kernel_size=parse_ksize(options["k"]),
+                exp_ratio=float(options["e"]),
+                se_ratio=float(options.get("se", 0.0)),
+                noskip=skip is False,
+            )
+        )
+    elif block_type == "uir":
+        start_kernel_size = parse_ksize(options.get("a", "0"))
+        end_kernel_size = parse_ksize(options.get("p", "0"))
+        block_args.update(
+            dict(
+                dw_kernel_size_start=start_kernel_size,
+                dw_kernel_size_mid=parse_ksize(options["k"]),
+                dw_kernel_size_end=end_kernel_size,
+                exp_ratio=float(options["e"]),
+                se_ratio=float(options.get("se", 0.0)),
+                noskip=skip is False,
+            )
+        )
+    elif block_type == "er":
+        block_args.update(
+            dict(
+                exp_kernel_size=parse_ksize(options["k"]),
+                pw_kernel_size=1,
+                exp_ratio=float(options["e"]),
+                se_ratio=float(options.get("se", 0.0)),
+                noskip=skip is False,
+            )
+        )
+    elif block_type in ("mqa", "mha"):
+        key_dim_val = int(options.get("d", "64"))
+        block_args.update(
+            dict(
+                num_heads=int(options.get("h", "12")),
+                key_dim=key_dim_val,
+                value_dim=key_dim_val,
+                kv_stride=int(options.get("v", "1")),
+                use_cpe=bool(int(options.get("cpe", "0"))),
+            )
+        )
+    return block_args, num_repeat
+
+
+def decode_arch_def(arch_def):
+    arch_args = []
+    for _, block_strings in enumerate(arch_def):
+        stack_args = []
+        for block_str in block_strings:
+            ba, rep = decode_block_str(block_str)
+            stack_args.extend([deepcopy(ba) for _ in range(rep)])
+        arch_args.append(stack_args)
+    return arch_args
+
+
+def convert_arch_def_to_stackwise(arch_def):
+    decoded_args = decode_arch_def(arch_def)
+    stackwise_params = {
+        k: []
+        for k in [
+            "stackwise_block_types",
+            "stackwise_num_blocks",
+            "stackwise_num_filters",
+            "stackwise_strides",
+            "stackwise_act_layers",
+            "stackwise_exp_ratios",
+            "stackwise_se_ratios",
+            "stackwise_dw_kernel_sizes",
+            "stackwise_dw_start_kernel_sizes",
+            "stackwise_dw_end_kernel_sizes",
+            "stackwise_exp_kernel_sizes",
+            "stackwise_pw_kernel_sizes",
+            "stackwise_num_heads",
+            "stackwise_key_dims",
+            "stackwise_value_dims",
+            "stackwise_kv_strides",
+            "stackwise_use_cpe",
+        ]
+    }
+    for stack in decoded_args:
+        stackwise_params["stackwise_num_blocks"].append(len(stack))
+        current_stack_params = {
+            k: [] for k in stackwise_params if k != "stackwise_num_blocks"
+        }
+        for block in stack:
+            current_stack_params["stackwise_block_types"].append(
+                block.get("block_type")
+            )
+            current_stack_params["stackwise_num_filters"].append(
+                block.get("out_chs")
+            )
+            current_stack_params["stackwise_strides"].append(
+                block.get("stride")
+            )
+            current_stack_params["stackwise_act_layers"].append(
+                block.get("act_layer")
+            )
+            current_stack_params["stackwise_exp_ratios"].append(
+                block.get("exp_ratio", 0.0)
+            )
+            current_stack_params["stackwise_se_ratios"].append(
+                block.get("se_ratio", 0.0)
+            )
+            current_stack_params["stackwise_dw_kernel_sizes"].append(
+                block.get("dw_kernel_size", block.get("dw_kernel_size_mid", 0))
+            )
+            current_stack_params["stackwise_dw_start_kernel_sizes"].append(
+                block.get("dw_kernel_size_start", 0)
+            )
+            current_stack_params["stackwise_dw_end_kernel_sizes"].append(
+                block.get("dw_kernel_size_end", 0)
+            )
+            current_stack_params["stackwise_exp_kernel_sizes"].append(
+                block.get("exp_kernel_size", 0)
+            )
+            current_stack_params["stackwise_pw_kernel_sizes"].append(
+                block.get("pw_kernel_size", 0)
+            )
+            current_stack_params["stackwise_num_heads"].append(
+                block.get("num_heads", 0)
+            )
+            current_stack_params["stackwise_key_dims"].append(
+                block.get("key_dim", 0)
+            )
+            current_stack_params["stackwise_value_dims"].append(
+                block.get("value_dim", 0)
+            )
+            current_stack_params["stackwise_kv_strides"].append(
+                block.get("kv_stride", 0)
+            )
+            current_stack_params["stackwise_use_cpe"].append(
+                block.get("use_cpe", False)
+            )
+        for key, value in current_stack_params.items():
+            stackwise_params[key].append(value)
+    return stackwise_params
+
+
+class MobileNetV5Builder:
+    """Builds a MobileNetV5 model from a decoded architecture definition.
+
+    This class takes a decoded architecture definition and constructs a list of
+    network stages, where each stage is a list of blocks. It handles channel
+    rounding, stride management, and feature extraction points.
+
+    Args:
+        output_stride: int. The desired output stride of the network.
+        pad_type: str. The padding type for convolutions.
+        round_chs_fn: callable. A function to round the number of channels.
+        se_from_exp: bool. If `True`, SE channel reduction is based on the
+            expanded channels.
+        act_layer: str. The default activation function for blocks.
+        norm_layer: str. The default normalization layer for blocks.
+        aa_layer: keras.layers.Layer. An optional anti-aliasing layer.
+        se_layer: keras.layers.Layer. The Squeeze-and-Excitation layer to use.
+        drop_path_rate: float. The stochastic depth rate for the network.
+        layer_scale_init_value: float. The initial value for layer scale.
+        feature_location: str. Where to extract features from, either
+            `"bottleneck"`, `"expansion"`, or `""`.
+        data_format: str. The format of the input data, either
+            `"channels_last"` or `"channels_first"`.
+        channel_axis: int. The axis representing the channels in the input
+            tensor.
+    """
+
+    def __init__(
+        self,
+        output_stride=32,
+        pad_type="same",
+        round_chs_fn=round_channels,
+        se_from_exp=False,
+        act_layer="relu",
+        norm_layer="batch_norm",
+        aa_layer=None,
+        se_layer=None,
+        drop_path_rate=0.0,
+        layer_scale_init_value=None,
+        feature_location="",
+        data_format=None,
+        channel_axis=None,
+        dtype=None,
+    ):
+        self.output_stride = output_stride
+        self.pad_type = pad_type
+        self.data_format = data_format
+        self.channel_axis = channel_axis
+        self.round_chs_fn = round_chs_fn
+        self.se_from_exp = se_from_exp
+        self.act_layer = act_layer
+        self.norm_layer = norm_layer
+        self.aa_layer = aa_layer
+        self.se_layer = se_layer
+        self.drop_path_rate = drop_path_rate
+        self.layer_scale_init_value = layer_scale_init_value
+        self.dtype = dtype
+        if feature_location == "depthwise":
+            feature_location = "expansion"
+        self.feature_location = feature_location
+        assert feature_location in ("bottleneck", "expansion", "")
+        self.in_chs = None
+        self.features = []
+
+    def _make_block(self, ba, block_idx, block_count):
+        drop_path_rate = self.drop_path_rate * block_idx / block_count
+        bt = ba.pop("block_type")
+        ba["filters"] = self.round_chs_fn(ba.pop("out_chs"))
+        s2d = ba.get("s2d", 0)
+        if s2d > 0:
+            ba["filters"] *= 4
+        if "expansion_in_chs" in ba and ba["expansion_in_chs"]:
+            ba["expansion_in_chs"] = self.round_chs_fn(ba["expansion_in_chs"])
+        ba["pad_type"] = self.pad_type
+        ba["act_layer"] = (
+            ba.get("act_layer")
+            if ba.get("act_layer") is not None
+            else self.act_layer
+        )
+        assert ba["act_layer"] is not None
+        ba["norm_layer"] = self.norm_layer
+        ba["drop_path_rate"] = drop_path_rate
+        ba["data_format"] = self.data_format
+        ba["channel_axis"] = self.channel_axis
+        ba["dtype"] = self.dtype
+        if bt in ("ir", "er", "uir", "ds", "dsa"):
+            se_ratio = ba.pop("se_ratio", None)
+            if se_ratio and self.se_layer is not None:
+                if not self.se_from_exp:
+                    se_ratio /= ba.get("exp_ratio", 1.0)
+                if s2d == 1:
+                    se_ratio /= 4
+                ba["se_layer"] = lambda channels: self.se_layer(
+                    filters=channels,
+                    bottleneck_filters=adjust_channels(channels * se_ratio),
+                    squeeze_activation=ba["act_layer"],
+                    excite_activation="sigmoid",
+                    data_format=self.data_format,
+                    dtype=self.dtype,
+                )
+            else:
+                ba["se_layer"] = None
+        ba.pop("aa_layer", None)
+        if bt == "ir":
+            padding = 0
+            if ba["pad_type"].lower() in ("", "same"):
+                kernel_size = ba["dw_kernel_size"]
+                if isinstance(kernel_size, (list, tuple)):
+                    kernel_size = kernel_size[0]
+                padding = (kernel_size - 1) // 2
+            block = (
+                CondConvResidual(**ba)
+                if ba.get("num_experts", 0) > 0
+                else InvertedResidualBlock(
+                    expansion=ba["exp_ratio"],
+                    infilters=self.in_chs,
+                    filters=ba["filters"],
+                    kernel_size=ba["dw_kernel_size"],
+                    stride=ba["stride"],
+                    padding=padding,
+                    squeeze_excite_ratio=ba.pop("se_ratio", None),
+                    activation=ba["act_layer"],
+                )
+            )
+        elif bt == "ds" or bt == "dsa":
+            block = DepthwiseConvBlock(
+                infilters=self.in_chs,
+                filters=ba["filters"],
+                kernel_size=ba["dw_kernel_size"],
+                stride=ba["stride"],
+                squeeze_excite_ratio=ba.pop("se_ratio", None),
+                residual=not ba["noskip"],
+                dtype=self.dtype,
+            )
+        elif bt == "er":
+            block = EdgeResidual(**ba)
+        elif bt == "cn":
+            block = ConvBnActBlock(out_chs=ba.pop("filters"), **ba)
+        elif bt == "uir":
+            block = UniversalInvertedResidual(
+                **ba, layer_scale_init_value=self.layer_scale_init_value
+            )
+        elif bt == "mqa":
+            ba.pop("act_layer", None)
+            block = MobileAttention(
+                **ba,
+                use_multi_query=True,
+                layer_scale_init_value=self.layer_scale_init_value,
+            )
+        elif bt == "mha":
+            ba.pop("act_layer", None)
+            block = MobileAttention(
+                **ba, layer_scale_init_value=self.layer_scale_init_value
+            )
+        else:
+            raise ValueError(f"Unknown block type ({bt}) while building model.")
+        self.in_chs = ba["filters"]
+        return block
+
+    def __call__(self, in_chs, model_block_args):
+        self.in_chs = in_chs
+        total_block_count = sum([len(x) for x in model_block_args])
+        total_block_idx = 0
+        current_stride = 2
+        current_dilation = 1
+        stages = []
+        if model_block_args[0][0]["stride"] > 1:
+            feature_info = dict(
+                module="conv_stem",
+                num_chs=in_chs,
+                stage=0,
+                reduction=current_stride,
+            )
+            self.features.append(feature_info)
+        space2depth = 0
+        for stack_idx, stack_args in enumerate(model_block_args):
+            blocks = []
+            for block_idx, block_args in enumerate(stack_args):
+                last_block = block_idx + 1 == len(stack_args)
+                in_chs_for_current_block = self.in_chs
+                assert block_args["stride"] in (1, 2)
+                if block_idx >= 1:
+                    block_args["stride"] = 1
+                if not space2depth and block_args.pop("s2d", False):
+                    assert block_args["stride"] == 1
+                    space2depth = 1
+                if space2depth > 0:
+                    if space2depth == 2 and block_args["stride"] == 2:
+                        block_args["stride"] = 1
+                        block_args["exp_ratio"] /= 4
+                        space2depth = 0
+                    else:
+                        block_args["s2d"] = space2depth
+                next_dilation = current_dilation
+                if block_args["stride"] > 1:
+                    next_output_stride = current_stride * block_args["stride"]
+                    if next_output_stride > self.output_stride:
+                        next_dilation = current_dilation * block_args["stride"]
+                        block_args["stride"] = 1
+                    else:
+                        current_stride = next_output_stride
+                block_args["dilation"] = current_dilation
+                if next_dilation != current_dilation:
+                    current_dilation = next_dilation
+                block = self._make_block(
+                    block_args.copy(), total_block_idx, total_block_count
+                )
+                blocks.append(block)
+                if space2depth == 1:
+                    space2depth = 2
+                extract_features = False
+                if last_block:
+                    next_stack_idx = stack_idx + 1
+                    extract_features = (
+                        next_stack_idx >= len(model_block_args)
+                        or model_block_args[next_stack_idx][0]["stride"] > 1
+                    )
+                if extract_features:
+                    num_chs = 0
+                    module_name = f"blocks.{stack_idx}.{block_idx}"
+                    if self.feature_location == "expansion":
+                        bt = block_args.get("block_type")
+                        if bt in ["ir", "er", "uir"]:
+                            exp_ratio = block_args.get("exp_ratio", 1.0)
+                            num_chs = self.round_chs_fn(
+                                in_chs_for_current_block * exp_ratio
+                            )
+                        else:
+                            num_chs = in_chs_for_current_block
+                    else:
+                        num_chs = self.in_chs
+                        module_name = f"blocks.{stack_idx}"
+
+                    feature_info = dict(
+                        stage=stack_idx + 1,
+                        reduction=current_stride,
+                        num_chs=num_chs,
+                        module=module_name,
+                    )
+                    self.features.append(feature_info)
+                total_block_idx += 1
+            stages.append(blocks)
+        return stages

keras_hub/src/models/mobilenetv5/mobilenetv5_image_classifier.py

@@ -0,0 +1,157 @@
+import keras
+
+from keras_hub.src.api_export import keras_hub_export
+from keras_hub.src.models.image_classifier import ImageClassifier
+from keras_hub.src.models.mobilenetv5.mobilenetv5_backbone import (
+    MobileNetV5Backbone,
+)
+from keras_hub.src.models.mobilenetv5.mobilenetv5_image_classifier_preprocessor import (  # noqa: E501
+    MobileNetV5ImageClassifierPreprocessor,
+)
+from keras_hub.src.models.mobilenetv5.mobilenetv5_layers import ConvNormAct
+from keras_hub.src.models.mobilenetv5.mobilenetv5_utils import (
+    SelectAdaptivePool2d,
+)
+from keras_hub.src.models.task import Task
+
+
+@keras_hub_export("keras_hub.models.MobileNetV5ImageClassifier")
+class MobileNetV5ImageClassifier(ImageClassifier):
+    """An end-to-end MobileNetV5 model for image classification.
+
+    This model attaches a classification head to a `MobileNetV5Backbone`.
+    The head consists of a global pooling layer, an optional convolutional
+    head, a dropout layer, and a final dense classifier layer.
+
+    This model can optionally be configured with a `preprocessor` layer, in
+    which case it will automatically apply preprocessing to image inputs during
+    `fit()`, `predict()`, and `evaluate()`.
+
+    Args:
+        backbone: A `keras_hub.models.MobileNetV5Backbone` instance.
+        num_classes: int. The number of classes for the classification head.
+        preprocessor: A `keras_hub.models.ImageClassifierPreprocessor` or
+            `None`. If `None`, this model will not apply preprocessing.
+        head_hidden_size: int. The number of channels in the convolutional
+            head.
+        global_pool: str. The type of global pooling to use.
+        drop_rate: float. The dropout rate for the head.
+        head_dtype: string or `keras.mixed_precision.DTypePolicy`. The dtype to
+            use for the head computations and weights.
+
+    Example:
+    ```python
+    import keras
+    from keras_hub.models import MobileNetV5Backbone
+    from keras_hub.models import MobileNetV5ImageClassifier
+
+    # Randomly initialized task model with a custom config.
+    model_args = {
+        "stackwise_block_types": [["er"], ["uir", "uir"]],
+        "stackwise_num_blocks": [1, 2],
+        "stackwise_num_filters": [[24], [48, 48]],
+        "stackwise_strides": [[2], [2, 1]],
+        "stackwise_act_layers": [["relu"], ["relu", "relu"]],
+        "stackwise_exp_ratios": [[4.0], [6.0, 6.0]],
+        "stackwise_se_ratios": [[0.0], [0.0, 0.0]],
+        "stackwise_dw_kernel_sizes": [[0], [5, 5]],
+        "stackwise_dw_start_kernel_sizes": [[0], [0, 0]],
+        "stackwise_dw_end_kernel_sizes": [[0], [0, 0]],
+        "stackwise_exp_kernel_sizes": [[3], [0, 0]],
+        "stackwise_pw_kernel_sizes": [[1], [0, 0]],
+        "stackwise_num_heads": [[0], [0, 0]],
+        "stackwise_key_dims": [[0], [0, 0]],
+        "stackwise_value_dims": [[0], [0, 0]],
+        "stackwise_kv_strides": [[0], [0, 0]],
+        "stackwise_use_cpe": [[False], [False, False]],
+        "use_msfa": False,
+    }
+    backbone = MobileNetV5Backbone(**model_args)
+    model = MobileNetV5ImageClassifier(backbone, 1000)
+    images = keras.ops.ones((1, 224, 224, 3))
+    output = model.predict(images)
+    ```
+    """
+
+    backbone_cls = MobileNetV5Backbone
+    preprocessor_cls = MobileNetV5ImageClassifierPreprocessor
+
+    def __init__(
+        self,
+        backbone,
+        num_classes,
+        preprocessor=None,
+        head_hidden_size=2048,
+        global_pool="avg",
+        drop_rate=0.0,
+        head_dtype=None,
+        **kwargs,
+    ):
+        head_dtype = head_dtype or backbone.dtype_policy
+        data_format = getattr(backbone, "data_format", "channels_last")
+
+        # === Layers ===
+        self.backbone = backbone
+        self.preprocessor = preprocessor
+        if backbone.use_msfa:
+            self.global_pool = SelectAdaptivePool2d(
+                pool_type=global_pool, data_format=data_format, flatten=True
+            )
+            self.conv_head = None
+            self.flatten = None
+        else:
+            self.global_pool = SelectAdaptivePool2d(
+                pool_type=global_pool, data_format=data_format, flatten=False
+            )
+            self.conv_head = ConvNormAct(
+                out_chs=head_hidden_size,
+                kernel_size=1,
+                pad_type="same",
+                norm_layer=backbone.norm_layer,
+                act_layer=backbone.act_layer,
+                bias=False,
+                name="conv_head",
+                dtype=head_dtype,
+            )
+            self.flatten = keras.layers.Flatten(dtype=head_dtype)
+        self.dropout = (
+            keras.layers.Dropout(drop_rate, dtype=head_dtype)
+            if drop_rate > 0.0
+            else None
+        )
+        self.classifier = (
+            keras.layers.Dense(num_classes, dtype=head_dtype, name="classifier")
+            if num_classes > 0
+            else keras.layers.Activation("linear", name="identity_classifier")
+        )
+
+        # === Functional Model ===
+        inputs = self.backbone.input
+        x = self.backbone(inputs)
+        x = self.global_pool(x)
+        if self.conv_head is not None:
+            x = self.conv_head(x)
+        if self.flatten is not None:
+            x = self.flatten(x)
+        if self.dropout is not None:
+            x = self.dropout(x)
+        outputs = self.classifier(x)
+        Task.__init__(self, inputs=inputs, outputs=outputs, **kwargs)
+
+        # === Config ===
+        self.num_classes = num_classes
+        self.head_hidden_size = head_hidden_size
+        self.global_pool_type = global_pool
+        self.drop_rate = drop_rate
+
+    def get_config(self):
+        config = Task.get_config(self)
+        config.update(
+            {
+                "num_classes": self.num_classes,
+                "head_hidden_size": self.head_hidden_size,
+                "global_pool": self.global_pool_type,
+                "drop_rate": self.drop_rate,
+            }
+        )
+        return config

keras_hub/src/models/mobilenetv5/mobilenetv5_image_classifier_preprocessor.py

@@ -0,0 +1,16 @@
+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.mobilenetv5.mobilenetv5_backbone import (
+    MobileNetV5Backbone,
+)
+from keras_hub.src.models.mobilenetv5.mobilenetv5_image_converter import (
+    MobileNetV5ImageConverter,
+)
+
+
+@keras_hub_export("keras_hub.models.MobileNetV5ImageClassifierPreprocessor")
+class MobileNetV5ImageClassifierPreprocessor(ImageClassifierPreprocessor):
+    backbone_cls = MobileNetV5Backbone
+    image_converter_cls = MobileNetV5ImageConverter

keras_hub/src/models/mobilenetv5/mobilenetv5_image_converter.py

@@ -0,0 +1,10 @@
+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.mobilenetv5.mobilenetv5_backbone import (
+    MobileNetV5Backbone,
+)
+
+
+@keras_hub_export("keras_hub.layers.MobileNetV5ImageConverter")
+class MobileNetV5ImageConverter(ImageConverter):
+    backbone_cls = MobileNetV5Backbone