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

keras_hub/src/models/vit/vit_image_classifier.py

@@ -0,0 +1,187 @@
+import keras
+from keras import ops
+
+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.vit.vit_backbone import ViTBackbone
+from keras_hub.src.models.vit.vit_image_classifier_preprocessor import (
+    ViTImageClassifierPreprocessor,
+)
+
+
+@keras_hub_export("keras_hub.models.ViTImageClassifier")
+class ViTImageClassifier(ImageClassifier):
+    """ViT image classification task.
+
+    `ViTImageClassifier` tasks wrap a `keras_hub.models.ViTBackbone` and
+    a `keras_hub.models.Preprocessor` to create a model that can be used for
+    image classification. `ViTImageClassifier` 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)`.
+
+    Not that unlike `keras_hub.model.ImageClassifier`, the `ViTImageClassifier`
+    we pluck out `cls_token` which is first seqence from the backbone.
+
+    Args:
+        backbone: A `keras_hub.models.ViTBackbone` 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: String specifying the classification strategy. The choice
+            impacts the dimensionality and nature of the feature vector used for
+            classification.
+                `"token"`:  A single vector (class token) representing the
+                    overall image features.
+                `"gap"`: A single vector representing the average features
+                    across the spatial dimensions.
+        intermediate_dim: Optional dimensionality of the intermediate
+            representation layer before the final classification layer.
+            If `None`, the output of the transformer is directly used.
+            Defaults to `None`.
+        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:
+
+    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.ViTImageClassifier.from_preset(
+        "vgg_16_imagenet"
+    )
+    classifier.predict(images)
+    ```
+
+    Call `fit()` on a single batch.
+    ```python
+    # Load preset and train
+    images = np.random.randint(0, 256, size=(2, 224, 224, 3))
+    labels = [0, 3]
+    classifier = keras_hub.models.VGGImageClassifier.from_preset(
+        "vit_base_patch16_224"
+    )
+    classifier.fit(x=images, y=labels, batch_size=2)
+    ```
+
+    Call `fit()` with custom loss, optimizer and backbone.
+    ```python
+    classifier = keras_hub.models.VGGImageClassifier.from_preset(
+        "vit_base_patch16_224"
+    )
+    classifier.compile(
+        loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
+        optimizer=keras.optimizers.Adam(5e-5),
+    )
+    classifier.backbone.trainable = False
+    classifier.fit(x=images, y=labels, batch_size=2)
+    ```
+
+    Custom backbone.
+    ```python
+    images = np.random.randint(0, 256, size=(2, 224, 224, 3))
+    labels = [0, 3]
+    model = keras_hub.models.ViTBackbone(
+        image_shape = (224, 224, 3),
+        patch_size=16,
+        num_layers=6,
+        num_heads=3,
+        hidden_dim=768,
+        mlp_dim=2048
+    )
+    classifier = keras_hub.models.ViTImageClassifier(
+        backbone=backbone,
+        num_classes=4,
+    )
+    classifier.fit(x=images, y=labels, batch_size=2)
+    ```
+    """
+
+    backbone_cls = ViTBackbone
+    preprocessor_cls = ViTImageClassifierPreprocessor
+
+    def __init__(
+        self,
+        backbone,
+        num_classes,
+        preprocessor=None,
+        pooling="token",
+        intermediate_dim=None,
+        activation=None,
+        dropout=0.0,
+        head_dtype=None,
+        **kwargs,
+    ):
+        head_dtype = head_dtype or backbone.dtype_policy
+
+        # === Layers ===
+        self.backbone = backbone
+        self.preprocessor = preprocessor
+
+        if intermediate_dim is not None:
+            self.intermediate_layer = keras.layers.Dense(
+                intermediate_dim, activation="tanh", name="pre_logits"
+            )
+
+        self.dropout = keras.layers.Dropout(
+            rate=dropout,
+            dtype=head_dtype,
+            name="output_dropout",
+        )
+        self.output_dense = keras.layers.Dense(
+            num_classes,
+            activation=activation,
+            dtype=head_dtype,
+            name="predictions",
+        )
+
+        # === Functional Model ===
+        inputs = self.backbone.input
+        x = self.backbone(inputs)
+        if pooling == "token":
+            x = x[:, 0]
+        elif pooling == "gap":
+            ndim = len(ops.shape(x))
+            x = ops.mean(x, axis=list(range(1, ndim - 1)))  # (1,) or (1,2)
+
+        if intermediate_dim is not None:
+            x = self.intermediate_layer(x)
+
+        x = self.dropout(x)
+        outputs = self.output_dense(x)
+
+        # Skip the parent class functional model.
+        Task.__init__(
+            self,
+            inputs=inputs,
+            outputs=outputs,
+            **kwargs,
+        )
+
+        # === config ===
+        self.num_classes = num_classes
+        self.pooling = pooling
+        self.intermediate_dim = intermediate_dim
+        self.activation = activation
+        self.dropout = dropout
+
+    def get_config(self):
+        # Backbone serialized in `super`
+        config = super().get_config()
+        config.update(
+            {
+                "num_classes": self.num_classes,
+                "pooling": self.pooling,
+                "intermediate_dim": self.intermediate_dim,
+                "activation": self.activation,
+                "dropout": self.dropout,
+            }
+        )
+        return config

keras_hub/src/models/vit/vit_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.vit.vit_backbone import ViTBackbone
+from keras_hub.src.models.vit.vit_image_converter import ViTImageConverter
+
+
+@keras_hub_export("keras_hub.models.ViTImageClassifierPreprocessor")
+class ViTImageClassifierPreprocessor(ImageClassifierPreprocessor):
+    backbone_cls = ViTBackbone
+    image_converter_cls = ViTImageConverter

keras_hub/src/models/vit/vit_image_converter.py

@@ -0,0 +1,73 @@
+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.vit.vit_backbone import ViTBackbone
+from keras_hub.src.utils.tensor_utils import preprocessing_function
+
+
+@keras_hub_export("keras_hub.layers.ViTImageConverter")
+class ViTImageConverter(ImageConverter):
+    """Converts images to the format expected by a ViT model.
+
+    This layer performs image normalization using mean and standard deviation
+    values. By default, it uses the same normalization as the
+    "google/vit-large-patch16-224" model on Hugging Face:
+    `norm_mean=[0.5, 0.5, 0.5]` and `norm_std=[0.5, 0.5, 0.5]`
+    ([reference](https://huggingface.co/google/vit-large-patch16-224/blob/main/preprocessor_config.json)).
+    These defaults are suitable for models pretrained using this normalization.
+
+    Args:
+        norm_mean: list or tuple of floats. Mean values for image normalization.
+            Defaults to `[0.5, 0.5, 0.5]`.
+        norm_std: list or tuple of floats. Standard deviation values for
+            image normalization. Defaults to `[0.5, 0.5, 0.5]`.
+        **kwargs: Additional keyword arguments passed to
+            `keras_hub.layers.preprocessing.ImageConverter`.
+
+    Examples:
+    ```python
+    import keras
+    import numpy as np
+    from keras_hub.src.layers import ViTImageConverter
+
+    # Example image (replace with your actual image data)
+    image = np.random.rand(1, 224, 224, 3)  # Example: (B, H, W, C)
+
+    # Create a ViTImageConverter instance
+    converter = ViTImageConverter(
+        image_size=(28,28),
+        scale=1/255.
+    )
+    # Preprocess the image
+    preprocessed_image = converter(image)
+    ```
+    """
+
+    backbone_cls = ViTBackbone
+
+    def __init__(
+        self, norm_mean=[0.5, 0.5, 0.5], norm_std=[0.5, 0.5, 0.5], **kwargs
+    ):
+        super().__init__(**kwargs)
+        self.norm_mean = norm_mean
+        self.norm_std = norm_std
+
+    @preprocessing_function
+    def call(self, inputs):
+        x = super().call(inputs)
+        # By default normalize using imagenet mean and std
+        if self.norm_mean:
+            x = x - self._expand_non_channel_dims(self.norm_mean, x)
+        if self.norm_std:
+            x = x / self._expand_non_channel_dims(self.norm_std, x)
+
+        return x
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "norm_mean": self.norm_mean,
+                "norm_std": self.norm_std,
+            }
+        )
+        return config

keras_hub/src/models/vit/vit_layers.py

@@ -0,0 +1,391 @@
+import keras
+from keras import ops
+
+from keras_hub.src.utils.keras_utils import standardize_data_format
+
+
+class MLP(keras.layers.Layer):
+    """Multi-Layer Perceptron (MLP) block.
+
+    Args:
+        hidden_dim: int. Dimensionality of the hidden representations.
+        mlp_dim: int. Dimensionality of the intermediate MLP layer.
+        use_bias: bool. Whether to use bias in the dense layers. Defaults to
+            `True`.
+        dropout_rate: float. Dropout rate. Between 0 and 1. Defaults to `0.0`.
+        **kwargs: Additional keyword arguments passed to `keras.layers.Layer`
+    """
+
+    def __init__(
+        self,
+        hidden_dim,
+        mlp_dim,
+        use_bias=True,
+        dropout_rate=0.0,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+
+        # === Config ===
+        self.hidden_dim = hidden_dim
+        self.mlp_dim = mlp_dim
+        self.use_bias = use_bias
+        self.dropout_rate = dropout_rate
+
+    def build(self, input_shape):
+        self.dense_1 = keras.layers.Dense(
+            units=self.mlp_dim,
+            use_bias=self.use_bias,
+            activation="gelu",
+            bias_initializer=(
+                keras.initializers.RandomNormal(stddev=1e-6)
+                if self.use_bias
+                else None
+            ),
+            dtype=self.dtype_policy,
+            name="dense_1",
+        )
+        self.dense_1.build(input_shape)
+        self.dense_2 = keras.layers.Dense(
+            units=self.hidden_dim,
+            use_bias=self.use_bias,
+            bias_initializer=(
+                keras.initializers.RandomNormal(stddev=1e-6)
+                if self.use_bias
+                else None
+            ),
+            dtype=self.dtype_policy,
+            name="dense_2",
+        )
+        self.dense_2.build((None, None, self.mlp_dim))
+        self.dropout = keras.layers.Dropout(
+            self.dropout_rate, dtype=self.dtype_policy, name="dropout"
+        )
+        self.built = True
+
+    def call(self, inputs):
+        x = self.dense_1(inputs)
+        x = self.dense_2(x)
+        out = self.dropout(x)
+        return out
+
+
+class ViTPatchingAndEmbedding(keras.layers.Layer):
+    """Patches the image and embeds the patches.
+
+    Args:
+        image_size: int. Size of the input image (height or width).
+            Assumed to be square.
+        patch_size: int. Size of each image patch.
+        hidden_dim: int. Dimensionality of the patch embeddings.
+        num_channels: int. Number of channels in the input image. Defaults to
+            `3`.
+        data_format: str. `"channels_last"` or `"channels_first"`. Defaults to
+            `None` (which uses `"channels_last"`).
+        **kwargs: Additional keyword arguments passed to `keras.layers.Layer`
+    """
+
+    def __init__(
+        self,
+        image_size,
+        patch_size,
+        hidden_dim,
+        num_channels=3,
+        data_format=None,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        num_patches = (image_size // patch_size) ** 2
+        num_positions = num_patches + 1
+
+        # === Config ===
+        self.image_size = image_size
+        self.patch_size = patch_size
+        self.hidden_dim = hidden_dim
+        self.num_channels = num_channels
+        self.num_patches = num_patches
+        self.num_positions = num_positions
+        self.data_format = standardize_data_format(data_format)
+
+    def build(self, input_shape):
+        self.class_token = self.add_weight(
+            shape=(
+                1,
+                1,
+                self.hidden_dim,
+            ),
+            initializer="random_normal",
+            dtype=self.variable_dtype,
+            name="class_token",
+        )
+        self.patch_embedding = keras.layers.Conv2D(
+            filters=self.hidden_dim,
+            kernel_size=self.patch_size,
+            strides=self.patch_size,
+            padding="valid",
+            activation=None,
+            dtype=self.dtype_policy,
+            data_format=self.data_format,
+            name="patch_embedding",
+        )
+        self.patch_embedding.build(input_shape)
+        self.position_embedding = keras.layers.Embedding(
+            self.num_positions,
+            self.hidden_dim,
+            dtype=self.dtype_policy,
+            embeddings_initializer=keras.initializers.RandomNormal(stddev=0.02),
+            name="position_embedding",
+        )
+        self.position_embedding.build((1, self.num_positions))
+        self.position_ids = keras.ops.expand_dims(
+            keras.ops.arange(self.num_positions), axis=0
+        )
+        self.built = True
+
+    def call(self, inputs):
+        patch_embeddings = self.patch_embedding(inputs)
+        if self.data_format == "channels_first":
+            patch_embeddings = ops.transpose(
+                patch_embeddings, axes=(0, 2, 3, 1)
+            )
+        embeddings_shape = ops.shape(patch_embeddings)
+        patch_embeddings = ops.reshape(
+            patch_embeddings, [embeddings_shape[0], -1, embeddings_shape[-1]]
+        )
+        class_token = ops.tile(self.class_token, (embeddings_shape[0], 1, 1))
+        position_embeddings = self.position_embedding(self.position_ids)
+        embeddings = ops.concatenate([class_token, patch_embeddings], axis=1)
+        return ops.add(embeddings, position_embeddings)
+
+    def compute_output_shape(self, input_shape):
+        return (
+            input_shape[0],
+            self.num_positions,
+            self.hidden_dim,
+        )
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "image_size": self.image_size,
+                "patch_size": self.patch_size,
+                "hidden_dim": self.hidden_dim,
+                "num_channels": self.num_channels,
+                "num_patches": self.num_patches,
+                "num_positions": self.num_positions,
+            }
+        )
+        return config
+
+
+class ViTEncoderBlock(keras.layers.Layer):
+    """Transformer encoder block.
+
+    Args:
+        num_heads: int. Number of attention heads.
+        hidden_dim: int. Dimensionality of the hidden representations.
+        mlp_dim: int. Dimensionality of the intermediate MLP layer.
+        use_mha_bias: bool. Whether to use bias in the multi-head attention
+            layer. Defaults to `True`.
+        use_mlp_bias: bool. Whether to use bias in the MLP layer. Defaults to
+            `True`.
+        dropout_rate: float. Dropout rate. Between 0 and 1. Defaults to `0.0`.
+        attention_dropout: float. Dropout rate for the attention mechanism.
+            Between 0 and 1. Defaults to `0.0`.
+        layer_norm_epsilon: float. Small float value for layer normalization
+            stability. Defaults to `1e-6`.
+        **kwargs: Additional keyword arguments passed to `keras.layers.Layer`
+    """
+
+    def __init__(
+        self,
+        num_heads,
+        hidden_dim,
+        mlp_dim,
+        use_mha_bias=True,
+        use_mlp_bias=True,
+        dropout_rate=0.0,
+        attention_dropout=0.0,
+        layer_norm_epsilon=1e-6,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+
+        key_dim = hidden_dim // num_heads
+
+        # === Config ===
+        self.num_heads = num_heads
+        self.hidden_dim = hidden_dim
+        self.key_dim = key_dim
+        self.mlp_dim = mlp_dim
+        self.use_mha_bias = use_mha_bias
+        self.use_mlp_bias = use_mlp_bias
+        self.dropout_rate = dropout_rate
+        self.attention_dropout = attention_dropout
+        self.layer_norm_epsilon = layer_norm_epsilon
+
+    def build(self, input_shape):
+        # Attention block
+        self.layer_norm_1 = keras.layers.LayerNormalization(
+            epsilon=self.layer_norm_epsilon,
+            name="ln_1",
+            dtype=self.dtype_policy,
+        )
+        self.layer_norm_1.build(input_shape)
+        self.mha = keras.layers.MultiHeadAttention(
+            num_heads=self.num_heads,
+            key_dim=self.key_dim,
+            use_bias=self.use_mha_bias,
+            dropout=self.attention_dropout,
+            name="mha",
+            dtype=self.dtype_policy,
+        )
+        self.mha.build(input_shape, input_shape)
+        self.dropout = keras.layers.Dropout(
+            self.dropout_rate, dtype=self.dtype_policy, name="dropout"
+        )
+
+        # MLP block
+        self.layer_norm_2 = keras.layers.LayerNormalization(
+            epsilon=self.layer_norm_epsilon,
+            name="ln_2",
+            dtype=self.dtype_policy,
+        )
+        self.layer_norm_2.build((None, None, self.hidden_dim))
+        self.mlp = MLP(
+            hidden_dim=self.hidden_dim,
+            mlp_dim=self.mlp_dim,
+            use_bias=self.use_mlp_bias,
+            name="mlp",
+            dtype=self.dtype_policy,
+        )
+        self.mlp.build((None, None, self.hidden_dim))
+        self.built = True
+
+    def call(self, inputs):
+        x = self.layer_norm_1(inputs)
+        x = self.mha(x, x)
+        x = self.dropout(x)
+        x = x + inputs
+
+        y = self.layer_norm_2(x)
+        y = self.mlp(y)
+
+        return x + y
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "num_heads": self.num_heads,
+                "hidden_dim": self.hidden_dim,
+                "key_dim": self.key_dim,
+                "mlp_dim": self.mlp_dim,
+                "use_mha_bias": self.use_mha_bias,
+                "use_mlp_bias": self.use_mlp_bias,
+                "dropout_rate": self.dropout_rate,
+                "attention_dropout": self.attention_dropout,
+                "layer_norm_epsilon": self.layer_norm_epsilon,
+            }
+        )
+        return config
+
+
+class ViTEncoder(keras.layers.Layer):
+    """Vision Transformer (ViT) encoder.
+
+    Args:
+        num_layers: int. Number of Transformer encoder blocks.
+        num_heads: int. Number of attention heads.
+        hidden_dim: int. Dimensionality of the hidden representations.
+        mlp_dim: int. Dimensionality of the intermediate MLP layer.
+        use_mha_bias: bool. Whether to use bias in the multi-head attention
+            layers. Defaults to `True`.
+        use_mlp_bias: bool. Whether to use bias in the MLP layers. Defaults to
+            `True`.
+        dropout_rate: float. Dropout rate. Between 0 and 1. Defaults to `0.0`.
+        attention_dropout: float. Dropout rate for the attention mechanism.
+            Between 0 and 1. Defaults to `0.0`.
+        layer_norm_epsilon: float. Small float value for layer normalization
+            tability. Defaults to `1e-6`.
+        **kwargs: Additional keyword arguments passed to `keras.layers.Layer`
+    """
+
+    def __init__(
+        self,
+        num_layers,
+        num_heads,
+        hidden_dim,
+        mlp_dim,
+        use_mha_bias=True,
+        use_mlp_bias=True,
+        dropout_rate=0.0,
+        attention_dropout=0.0,
+        layer_norm_epsilon=1e-6,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+
+        # === config ===
+        self.num_layers = num_layers
+        self.num_heads = num_heads
+        self.hidden_dim = hidden_dim
+        self.mlp_dim = mlp_dim
+        self.use_mha_bias = use_mha_bias
+        self.use_mlp_bias = use_mlp_bias
+        self.dropout_rate = dropout_rate
+        self.attention_dropout = attention_dropout
+        self.layer_norm_epsilon = layer_norm_epsilon
+
+    def build(self, input_shape):
+        self.encoder_layers = []
+        for i in range(self.num_layers):
+            encoder_block = ViTEncoderBlock(
+                num_heads=self.num_heads,
+                hidden_dim=self.hidden_dim,
+                mlp_dim=self.mlp_dim,
+                dropout_rate=self.dropout_rate,
+                use_mha_bias=self.use_mha_bias,
+                use_mlp_bias=self.use_mlp_bias,
+                attention_dropout=self.attention_dropout,
+                layer_norm_epsilon=self.layer_norm_epsilon,
+                dtype=self.dtype_policy,
+                name=f"tranformer_block_{i + 1}",
+            )
+            encoder_block.build((None, None, self.hidden_dim))
+            self.encoder_layers.append(encoder_block)
+        self.dropout = keras.layers.Dropout(
+            self.dropout_rate, dtype=self.dtype_policy, name="dropout"
+        )
+        self.layer_norm = keras.layers.LayerNormalization(
+            epsilon=self.layer_norm_epsilon,
+            dtype=self.dtype_policy,
+            name="ln",
+        )
+        self.layer_norm.build((None, None, self.hidden_dim))
+        self.built = True
+
+    def call(self, inputs):
+        x = self.dropout(inputs)
+        for i in range(self.num_layers):
+            x = self.encoder_layers[i](x)
+        x = self.layer_norm(x)
+        return x
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "num_layers": self.num_layers,
+                "num_heads": self.num_heads,
+                "hidden_dim": self.hidden_dim,
+                "mlp_dim": self.mlp_dim,
+                "use_mha_bias": self.use_mha_bias,
+                "use_mlp_bias": self.use_mlp_bias,
+                "dropout_rate": self.dropout_rate,
+                "attention_dropout": self.attention_dropout,
+                "layer_norm_epsilon": self.layer_norm_epsilon,
+            }
+        )
+        return config