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

keras_hub/src/models/vgg/__init__.py

@@ -0,0 +1,5 @@
+from keras_hub.src.models.vgg.vgg_backbone import VGGBackbone
+from keras_hub.src.models.vgg.vgg_presets import backbone_presets
+from keras_hub.src.utils.preset_utils import register_presets
+
+register_presets(backbone_presets, VGGBackbone)

keras_hub/src/models/vgg/vgg_backbone.py

@@ -20,25 +20,14 @@ class VGGBackbone(Backbone):
       stackwise_num_filters: list of ints, filter size for convolutional
             blocks per VGG block. For both VGG16 and VGG19 this is [
             64, 128, 256, 512, 512].
-      image_shape: tuple, optional shape tuple, defaults to (224, 224, 3).
-      pooling: bool, Optional pooling mode for feature extraction
-        when `include_top` is `False`.
-        - `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.
+      image_shape: tuple, optional shape tuple, defaults to (None, None, 3).
 
     Examples:
     ```python
     input_data = np.ones((2, 224, 224, 3), dtype="float32")
 
     # Pretrained VGG backbone.
-    model = keras_hub.models.VGGBackbone.from_preset("vgg16")
+    model = keras_hub.models.VGGBackbone.from_preset("vgg_16_imagenet")
     model(input_data)
 
     # Randomly initialized VGG backbone with a custom config.
@@ -46,7 +35,6 @@ class VGGBackbone(Backbone):
         stackwise_num_repeats = [2, 2, 3, 3, 3],
         stackwise_num_filters = [64, 128, 256, 512, 512],
         image_shape = (224, 224, 3),
-        pooling = "avg",
     )
     model(input_data)
     ```
@@ -56,16 +44,14 @@ class VGGBackbone(Backbone):
         self,
         stackwise_num_repeats,
         stackwise_num_filters,
-        image_shape=(224, 224, 3),
-        pooling="avg",
+        image_shape=(None, None, 3),
         **kwargs,
     ):
-
         # === Functional Model ===
         img_input = keras.layers.Input(shape=image_shape)
         x = img_input
 
-        for stack_index in range(len(stackwise_num_repeats) - 1):
+        for stack_index in range(len(stackwise_num_repeats)):
             x = apply_vgg_block(
                 x=x,
                 num_layers=stackwise_num_repeats[stack_index],
@@ -76,10 +62,6 @@ class VGGBackbone(Backbone):
                 max_pool=True,
                 name=f"block{stack_index + 1}",
             )
-        if pooling == "avg":
-            x = layers.GlobalAveragePooling2D()(x)
-        elif pooling == "max":
-            x = layers.GlobalMaxPooling2D()(x)
 
         super().__init__(inputs=img_input, outputs=x, **kwargs)
 
@@ -87,14 +69,12 @@ class VGGBackbone(Backbone):
         self.stackwise_num_repeats = stackwise_num_repeats
         self.stackwise_num_filters = stackwise_num_filters
         self.image_shape = image_shape
-        self.pooling = pooling
 
     def get_config(self):
         return {
             "stackwise_num_repeats": self.stackwise_num_repeats,
             "stackwise_num_filters": self.stackwise_num_filters,
             "image_shape": self.image_shape,
-            "pooling": self.pooling,
         }
 
 

keras_hub/src/models/vgg/vgg_image_classifier.py

@@ -2,58 +2,93 @@ 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.task import Task
 from keras_hub.src.models.vgg.vgg_backbone import VGGBackbone
+from keras_hub.src.models.vgg.vgg_image_classifier_preprocessor import (
+    VGGImageClassifierPreprocessor,
+)
 
 
 @keras_hub_export("keras_hub.models.VGGImageClassifier")
 class VGGImageClassifier(ImageClassifier):
-    """VGG16 image classifier task model.
+    """VGG image classification task.
 
-    Args:
-      backbone: A `keras_hub.models.VGGBackbone` instance.
-      num_classes: int, number of classes to predict.
-      pooling: str, type of pooling layer. Must be one of "avg", "max".
-      activation: Optional `str` or callable, defaults to "softmax". The
-          activation function to use on the Dense layer. Set `activation=None`
-          to return the output logits.
+    `VGGImageClassifier` tasks wrap a `keras_hub.models.VGGBackbone` and
+    a `keras_hub.models.Preprocessor` to create a model that can be used for
+    image classification. `VGGImageClassifier` tasks take an additional
+    `num_classes` argument, controlling the number of predicted output classes.
 
     To fine-tune with `fit()`, pass a dataset containing tuples of `(x, y)`
     labels where `x` is a string and `y` is a integer from `[0, num_classes)`.
-    All `ImageClassifier` tasks include a `from_preset()` constructor which can be
-    used to load a pre-trained config and weights.
+
+    Not that unlike `keras_hub.model.ImageClassifier`, the `VGGImageClassifier`
+    allows and defaults to `pooling="flatten"`, when inputs are flatten and
+    passed through two intermediate dense layers before the final output
+    projection.
+
+    Args:
+        backbone: A `keras_hub.models.VGGBackbone` instance or a `keras.Model`.
+        num_classes: int. The number of classes to predict.
+        preprocessor: `None`, a `keras_hub.models.Preprocessor` instance,
+            a `keras.Layer` instance, or a callable. If `None` no preprocessing
+            will be applied to the inputs.
+        pooling: `"flatten"`, `"avg"`, or `"max"`. The type of pooling to apply
+            on backbone output. The default is flatten to match the original
+            VGG implementation, where backbone inputs will be flattened and
+            passed through two dense layers with a `"relu"` activation.
+        pooling_hidden_dim: the output feature size of the pooling dense layers.
+            This only applies when `pooling="flatten"`.
+        activation: `None`, str, or callable. The activation function to use on
+            the `Dense` layer. Set `activation=None` to return the output
+            logits. Defaults to `"softmax"`.
+        head_dtype: `None`, str, or `keras.mixed_precision.DTypePolicy`. The
+            dtype to use for the classification head's computations and weights.
+
 
     Examples:
-    Train from preset
+
+    Call `predict()` to run inference.
+    ```python
+    # Load preset and train
+    images = np.random.randint(0, 256, size=(2, 224, 224, 3))
+    classifier = keras_hub.models.VGGImageClassifier.from_preset(
+        "vgg_16_imagenet"
+    )
+    classifier.predict(images)
+    ```
+
+    Call `fit()` on a single batch.
     ```python
     # Load preset and train
-    images = np.ones((2, 224, 224, 3), dtype="float32")
+    images = np.random.randint(0, 256, size=(2, 224, 224, 3))
     labels = [0, 3]
     classifier = keras_hub.models.VGGImageClassifier.from_preset(
-        'vgg_16_image_classifier')
+        "vgg_16_imagenet"
+    )
     classifier.fit(x=images, y=labels, batch_size=2)
+    ```
 
-    # Re-compile (e.g., with a new learning rate).
+    Call `fit()` with custom loss, optimizer and backbone.
+    ```python
+    classifier = keras_hub.models.VGGImageClassifier.from_preset(
+        "vgg_16_imagenet"
+    )
     classifier.compile(
         loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
         optimizer=keras.optimizers.Adam(5e-5),
-        jit_compile=True,
     )
-
-    # Access backbone programmatically (e.g., to change `trainable`).
     classifier.backbone.trainable = False
-    # Fit again.
     classifier.fit(x=images, y=labels, batch_size=2)
     ```
-    Custom backbone
+
+    Custom backbone.
     ```python
-    images = np.ones((2, 224, 224, 3), dtype="float32")
+    images = np.random.randint(0, 256, size=(2, 224, 224, 3))
     labels = [0, 3]
-
-    backbone = keras_hub.models.VGGBackbone(
+    model = keras_hub.models.VGGBackbone(
         stackwise_num_repeats = [2, 2, 3, 3, 3],
         stackwise_num_filters = [64, 128, 256, 512, 512],
         image_shape = (224, 224, 3),
-        pooling = "avg",
     )
     classifier = keras_hub.models.VGGImageClassifier(
         backbone=backbone,
@@ -64,31 +99,95 @@ class VGGImageClassifier(ImageClassifier):
     """
 
     backbone_cls = VGGBackbone
+    preprocessor_cls = VGGImageClassifierPreprocessor
 
     def __init__(
         self,
         backbone,
         num_classes,
-        activation="softmax",
-        preprocessor=None,  # adding this dummy arg for saved model test
-        # TODO: once preprocessor flow is figured out, this needs to be updated
+        preprocessor=None,
+        pooling="avg",
+        pooling_hidden_dim=4096,
+        activation=None,
+        dropout=0.0,
+        head_dtype=None,
         **kwargs,
     ):
+        head_dtype = head_dtype or backbone.dtype_policy
+        data_format = getattr(backbone, "data_format", None)
+
         # === Layers ===
         self.backbone = backbone
-        self.output_dense = keras.layers.Dense(
-            num_classes,
-            activation=activation,
-            name="predictions",
+        self.preprocessor = preprocessor
+        if pooling == "avg":
+            self.pooler = keras.layers.GlobalAveragePooling2D(
+                data_format,
+                dtype=head_dtype,
+                name="pooler",
+            )
+        elif pooling == "max":
+            self.pooler = keras.layers.GlobalMaxPooling2D(
+                data_format,
+                dtype=head_dtype,
+                name="pooler",
+            )
+        elif pooling == "flatten":
+            self.pooler = keras.Sequential(
+                [
+                    keras.layers.Flatten(name="flatten"),
+                    keras.layers.Dense(pooling_hidden_dim, activation="relu"),
+                    keras.layers.Dense(pooling_hidden_dim, activation="relu"),
+                ],
+                name="pooler",
+            )
+        else:
+            raise ValueError(
+                "Unknown `pooling` type. Polling should be either `'avg'` or "
+                f"`'max'`. Received: pooling={pooling}."
+            )
+
+        self.head = keras.Sequential(
+            [
+                keras.layers.Conv2D(
+                    filters=4096,
+                    kernel_size=7,
+                    name="fc1",
+                    activation=activation,
+                    use_bias=True,
+                    padding="same",
+                ),
+                keras.layers.Dropout(
+                    rate=dropout,
+                    dtype=head_dtype,
+                    name="output_dropout",
+                ),
+                keras.layers.Conv2D(
+                    filters=4096,
+                    kernel_size=1,
+                    name="fc2",
+                    activation=activation,
+                    use_bias=True,
+                    padding="same",
+                ),
+                self.pooler,
+                keras.layers.Dense(
+                    num_classes,
+                    activation=activation,
+                    dtype=head_dtype,
+                    name="predictions",
+                ),
+            ],
+            name="head",
         )
 
         # === Functional Model ===
         inputs = self.backbone.input
         x = self.backbone(inputs)
-        outputs = self.output_dense(x)
+        outputs = self.head(x)
 
-        # Instantiate using Functional API Model constructor
-        super().__init__(
+        # Skip the parent class functional model.
+        Task.__init__(
+            self,
             inputs=inputs,
             outputs=outputs,
             **kwargs,
@@ -97,6 +196,10 @@ class VGGImageClassifier(ImageClassifier):
         # === Config ===
         self.num_classes = num_classes
         self.activation = activation
+        self.pooling = pooling
+        self.pooling_hidden_dim = pooling_hidden_dim
+        self.dropout = dropout
+        self.preprocessor = preprocessor
 
     def get_config(self):
         # Backbone serialized in `super`
@@ -104,7 +207,10 @@ class VGGImageClassifier(ImageClassifier):
         config.update(
             {
                 "num_classes": self.num_classes,
+                "pooling": self.pooling,
                 "activation": self.activation,
+                "pooling_hidden_dim": self.pooling_hidden_dim,
+                "dropout": self.dropout,
             }
         )
         return config

keras_hub/src/models/vgg/vgg_image_classifier_preprocessor.py

@@ -0,0 +1,12 @@
+from keras_hub.src.api_export import keras_hub_export
+from keras_hub.src.models.image_classifier_preprocessor import (
+    ImageClassifierPreprocessor,
+)
+from keras_hub.src.models.vgg.vgg_backbone import VGGBackbone
+from keras_hub.src.models.vgg.vgg_image_converter import VGGImageConverter
+
+
+@keras_hub_export("keras_hub.models.VGGImageClassifierPreprocessor")
+class VGGImageClassifierPreprocessor(ImageClassifierPreprocessor):
+    backbone_cls = VGGBackbone
+    image_converter_cls = VGGImageConverter

keras_hub/src/models/vgg/vgg_image_converter.py

@@ -0,0 +1,8 @@
+from keras_hub.src.api_export import keras_hub_export
+from keras_hub.src.layers.preprocessing.image_converter import ImageConverter
+from keras_hub.src.models.vgg.vgg_backbone import VGGBackbone
+
+
+@keras_hub_export("keras_hub.layers.VGGImageConverter")
+class VGGImageConverter(ImageConverter):
+    backbone_cls = VGGBackbone

keras_hub/src/models/vgg/vgg_presets.py

@@ -0,0 +1,48 @@
+"""vgg preset configurations."""
+
+backbone_presets = {
+    "vgg_11_imagenet": {
+        "metadata": {
+            "description": (
+                "11-layer vgg model pre-trained on the ImageNet 1k dataset "
+                "at a 224x224 resolution."
+            ),
+            "params": 9220480,
+            "path": "vgg",
+        },
+        "kaggle_handle": "kaggle://keras/vgg/keras/vgg_11_imagenet/2",
+    },
+    "vgg_13_imagenet": {
+        "metadata": {
+            "description": (
+                "13-layer vgg model pre-trained on the ImageNet 1k dataset "
+                "at a 224x224 resolution."
+            ),
+            "params": 9404992,
+            "path": "vgg",
+        },
+        "kaggle_handle": "kaggle://keras/vgg/keras/vgg_13_imagenet/2",
+    },
+    "vgg_16_imagenet": {
+        "metadata": {
+            "description": (
+                "16-layer vgg model pre-trained on the ImageNet 1k dataset "
+                "at a 224x224 resolution."
+            ),
+            "params": 14714688,
+            "path": "vgg",
+        },
+        "kaggle_handle": "kaggle://keras/vgg/keras/vgg_16_imagenet/2",
+    },
+    "vgg_19_imagenet": {
+        "metadata": {
+            "description": (
+                "19-layer vgg model pre-trained on the ImageNet 1k dataset "
+                "at a 224x224 resolution."
+            ),
+            "params": 20024384,
+            "path": "vgg",
+        },
+        "kaggle_handle": "kaggle://keras/vgg/keras/vgg_19_imagenet/2",
+    },
+}

keras_hub/src/models/vit/__init__.py

@@ -0,0 +1,5 @@
+from keras_hub.src.models.vit.vit_backbone import ViTBackbone
+from keras_hub.src.models.vit.vit_presets import backbone_presets
+from keras_hub.src.utils.preset_utils import register_presets
+
+register_presets(backbone_presets, ViTBackbone)

keras_hub/src/models/vit/vit_backbone.py

@@ -0,0 +1,152 @@
+import keras
+
+from keras_hub.src.api_export import keras_hub_export
+from keras_hub.src.models.backbone import Backbone
+from keras_hub.src.models.vit.vit_layers import ViTEncoder
+from keras_hub.src.models.vit.vit_layers import ViTPatchingAndEmbedding
+from keras_hub.src.utils.keras_utils import standardize_data_format
+
+
+@keras_hub_export("keras_hub.models.ViTBackbone")
+class ViTBackbone(Backbone):
+    """Vision Transformer (ViT) backbone.
+
+    This backbone implements the Vision Transformer architecture as described in
+    [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://arxiv.org/abs/2010.11929).
+    It transforms the input image into a sequence of patches, embeds them, and
+    then processes them through a series of Transformer encoder layers.
+
+    Args:
+        image_shape: A tuple or list of 3 integers representing the shape of the
+            input image `(height, width, channels)`, `height` and `width` must
+            be equal.
+        patch_size: int. The size of each image patch, the input image will be
+            divided into patches of shape `(patch_size, patch_size)`.
+        num_layers: int. The number of transformer encoder layers.
+        num_heads: int. specifying the number of attention heads in each
+            Transformer encoder layer.
+        hidden_dim: int. The dimensionality of the hidden representations.
+        mlp_dim: int. The dimensionality of the intermediate MLP layer in
+            each Transformer encoder layer.
+        dropout_rate: float. The dropout rate for the Transformer encoder
+            layers.
+        attention_dropout: float. The dropout rate for the attention mechanism
+            in each Transformer encoder layer.
+        layer_norm_epsilon: float. Value used for numerical stability in
+            layer normalization.
+        use_mha_bias: bool. Whether to use bias in the multi-head
+            attention layers.
+        use_mlp_bias: bool. Whether to use bias in the MLP layers.
+        data_format: str.  `"channels_last"` or `"channels_first"`, specifying
+            the data format for the input image. If `None`, defaults to
+            `"channels_last"`.
+        dtype: The dtype of the layer weights. Defaults to None.
+        **kwargs: Additional keyword arguments to be passed to the parent
+            `Backbone` class.
+    """
+
+    def __init__(
+        self,
+        image_shape,
+        patch_size,
+        num_layers,
+        num_heads,
+        hidden_dim,
+        mlp_dim,
+        dropout_rate=0.0,
+        attention_dropout=0.0,
+        layer_norm_epsilon=1e-6,
+        use_mha_bias=True,
+        use_mlp_bias=True,
+        data_format=None,
+        dtype=None,
+        **kwargs,
+    ):
+        # === Laters ===
+        data_format = standardize_data_format(data_format)
+        h_axis, w_axis, channels_axis = (
+            (-3, -2, -1) if data_format == "channels_last" else (-2, -1, -3)
+        )
+        # Check that the input image is well specified.
+        if image_shape[h_axis] is None or image_shape[w_axis] is None:
+            raise ValueError(
+                f"Image shape must have defined height and width. Found `None` "
+                f"at index {h_axis} (height) or {w_axis} (width). "
+                f"Image shape: {image_shape}"
+            )
+        if image_shape[h_axis] != image_shape[w_axis]:
+            raise ValueError(
+                f"Image height and width must be equal. Found height: "
+                f"{image_shape[h_axis]}, width: {image_shape[w_axis]} at "
+                f"indices {h_axis} and {w_axis} respectively. Image shape: "
+                f"{image_shape}"
+            )
+
+        num_channels = image_shape[channels_axis]
+
+        # === Functional Model ===
+        inputs = keras.layers.Input(shape=image_shape)
+
+        x = ViTPatchingAndEmbedding(
+            image_size=image_shape[h_axis],
+            patch_size=patch_size,
+            hidden_dim=hidden_dim,
+            num_channels=num_channels,
+            data_format=data_format,
+            dtype=dtype,
+            name="vit_patching_and_embedding",
+        )(inputs)
+
+        output = ViTEncoder(
+            num_layers=num_layers,
+            num_heads=num_heads,
+            hidden_dim=hidden_dim,
+            mlp_dim=mlp_dim,
+            dropout_rate=dropout_rate,
+            attention_dropout=attention_dropout,
+            layer_norm_epsilon=layer_norm_epsilon,
+            use_mha_bias=use_mha_bias,
+            use_mlp_bias=use_mlp_bias,
+            dtype=dtype,
+            name="vit_encoder",
+        )(x)
+
+        super().__init__(
+            inputs=inputs,
+            outputs=output,
+            dtype=dtype,
+            **kwargs,
+        )
+
+        # === Config ===
+        self.image_shape = image_shape
+        self.patch_size = patch_size
+        self.num_layers = num_layers
+        self.num_heads = num_heads
+        self.hidden_dim = hidden_dim
+        self.mlp_dim = mlp_dim
+        self.dropout_rate = dropout_rate
+        self.attention_dropout = attention_dropout
+        self.layer_norm_epsilon = layer_norm_epsilon
+        self.use_mha_bias = use_mha_bias
+        self.use_mlp_bias = use_mlp_bias
+        self.data_format = data_format
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "image_shape": self.image_shape,
+                "patch_size": self.patch_size,
+                "num_layers": self.num_layers,
+                "num_heads": self.num_heads,
+                "hidden_dim": self.hidden_dim,
+                "mlp_dim": self.mlp_dim,
+                "dropout_rate": self.dropout_rate,
+                "attention_dropout": self.attention_dropout,
+                "layer_norm_epsilon": self.layer_norm_epsilon,
+                "use_mha_bias": self.use_mha_bias,
+                "use_mlp_bias": self.use_mlp_bias,
+            }
+        )
+        return config