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

keras_hub/src/models/gemma/gemma_backbone.py

@@ -44,10 +44,10 @@ class GemmaBackbone(Backbone):
             `hidden_dim / num_query_heads`. Defaults to True.
         use_post_ffw_norm: boolean. Whether to normalize after the feedforward
             block. Defaults to False.
-        use_post_attention_norm: boolean. Whether to normalize after the attention
-            block. Defaults to False.
-        attention_logit_soft_cap: None or int. Soft cap for the attention logits.
-            Defaults to None.
+        use_post_attention_norm: boolean. Whether to normalize after the
+            attention block. Defaults to False.
+        attention_logit_soft_cap: None or int. Soft cap for the attention
+            logits. Defaults to None.
         final_logit_soft_cap: None or int. Soft cap for the final logits.
             Defaults to None.
         use_sliding_window_attention boolean. Whether to use sliding local
@@ -205,7 +205,9 @@ class GemmaBackbone(Backbone):
                 "final_logit_soft_cap": self.final_logit_soft_cap,
                 "attention_logit_soft_cap": self.attention_logit_soft_cap,
                 "sliding_window_size": self.sliding_window_size,
-                "use_sliding_window_attention": self.use_sliding_window_attention,
+                "use_sliding_window_attention": (
+                    self.use_sliding_window_attention
+                ),
             }
         )
         return config
@@ -224,7 +226,8 @@ class GemmaBackbone(Backbone):
 
         Example:
         ```
-        # Feel free to change the mesh shape to balance data and model parallel
+        # Feel free to change the mesh shape to balance data and model
+        # parallelism
         mesh = keras.distribution.DeviceMesh(
             shape=(1, 8), axis_names=('batch', 'model'),
             devices=keras.distribution.list_devices())
@@ -232,11 +235,23 @@ class GemmaBackbone(Backbone):
             mesh, model_parallel_dim_name="model")
 
         distribution = keras.distribution.ModelParallel(
-            mesh, layout_map, batch_dim_name='batch')
+            layout_map=layout_map, batch_dim_name='batch')
         with distribution.scope():
            gemma_model = keras_hub.models.GemmaCausalLM.from_preset()
         ```
 
+        To see how the layout map was applied, load the model then run (for one
+        decoder block):
+        ```
+        embedding_layer = gemma_model.backbone.get_layer("token_embedding")
+        decoder_block_1 = gemma_model.backbone.get_layer('decoder_block_1')
+        for variable in embedding_layer.weights + decoder_block_1.weights:
+            print(
+                f'{variable.path:<58}  {str(variable.shape):<16} '
+                f'{str(variable.value.sharding.spec)}'
+            )
+        ```
+
         Args:
             device_mesh: The `keras.distribution.DeviceMesh` instance for
                 distribution.
@@ -246,25 +261,25 @@ class GemmaBackbone(Backbone):
                 the data should be partition on.
         Return:
             `keras.distribution.LayoutMap` that contains the sharding spec
-            of all the model weights.
+            for all the model weights.
         """
         # The weight path and shape of the Gemma backbone is like below (for 2G)
-        # token_embedding/embeddings,  (256128, 2048), 524550144
+        # token_embedding/embeddings,  (256128, 2048)
         # repeat block for decoder
         # ...
-        # decoder_block_17/pre_attention_norm/scale,  (2048,), 2048
-        # decoder_block_17/attention/query/kernel,  (8, 2048, 256), 4194304
-        # decoder_block_17/attention/key/kernel,  (8, 2048, 256), 4194304
-        # decoder_block_17/attention/value/kernel,  (8, 2048, 256), 4194304
-        # decoder_block_17/attention/attention_output/kernel,  (8, 256, 2048), 4194304
-        # decoder_block_17/pre_ffw_norm/scale,  (2048,), 2048
-        # decoder_block_17/ffw_gating/kernel,  (2048, 16384), 33554432
-        # decoder_block_17/ffw_gating_2/kernel,  (2048, 16384), 33554432
-        # decoder_block_17/ffw_linear/kernel,  (16384, 2048), 33554432
+        # decoder_block_17/pre_attention_norm/scale,  (2048,)
+        # decoder_block_17/attention/query/kernel,  (8, 2048, 256)
+        # decoder_block_17/attention/key/kernel,  (8, 2048, 256)
+        # decoder_block_17/attention/value/kernel,  (8, 2048, 256)
+        # decoder_block_17/attention/attention_output/kernel,  (8, 256, 2048)
+        # decoder_block_17/pre_ffw_norm/scale,  (2048,)
+        # decoder_block_17/ffw_gating/kernel,  (2048, 16384)
+        # decoder_block_17/ffw_gating_2/kernel,  (2048, 16384)
+        # decoder_block_17/ffw_linear/kernel,  (16384, 2048)
         if not isinstance(device_mesh, keras.distribution.DeviceMesh):
             raise ValueError(
-                "Invalid device_mesh type. Expected `keras.distribution.Device`,"
-                f" got {type(device_mesh)}"
+                "Invalid device_mesh type. Expected "
+                f"`keras.distribution.Device`, got {type(device_mesh)}"
             )
         if model_parallel_dim_name not in device_mesh.axis_names:
             raise ValueError(

keras_hub/src/models/gemma/gemma_causal_lm.py

@@ -187,8 +187,8 @@ class GemmaCausalLM(CausalLM):
         Args:
             token_ids: a dense int Tensor with shape `(batch_size, max_length)`.
             cache: a dense float Tensor, the cache of key and value.
-            cache_update_index: int, or int Tensor. The index of current inputs in the
-                whole sequence.
+            cache_update_index: int, or int Tensor. The index of current inputs
+                in the whole sequence.
 
         Returns:
             A (logits, hidden_states, cache) tuple. Where `logits` is the

keras_hub/src/models/gemma/gemma_decoder_block.py

@@ -220,7 +220,9 @@ class GemmaDecoderBlock(keras.layers.Layer):
                 "use_post_ffw_norm": self.use_post_ffw_norm,
                 "use_post_attention_norm": self.use_post_attention_norm,
                 "logit_soft_cap": self.logit_soft_cap,
-                "use_sliding_window_attention": self.use_sliding_window_attention,
+                "use_sliding_window_attention": (
+                    self.use_sliding_window_attention
+                ),
                 "sliding_window_size": self.sliding_window_size,
                 "query_head_dim_normalize": self.query_head_dim_normalize,
             }

keras_hub/src/models/gemma/gemma_presets.py

@@ -6,11 +6,9 @@ backbone_presets = {
         "metadata": {
             "description": "2 billion parameter, 18-layer, base Gemma model.",
             "params": 2506172416,
-            "official_name": "Gemma",
             "path": "gemma",
-            "model_card": "https://www.kaggle.com/models/google/gemma",
         },
-        "kaggle_handle": "kaggle://keras/gemma/keras/gemma_2b_en/2",
+        "kaggle_handle": "kaggle://keras/gemma/keras/gemma_2b_en/3",
     },
     "gemma_instruct_2b_en": {
         "metadata": {
@@ -18,11 +16,9 @@ backbone_presets = {
                 "2 billion parameter, 18-layer, instruction tuned Gemma model."
             ),
             "params": 2506172416,
-            "official_name": "Gemma",
             "path": "gemma",
-            "model_card": "https://www.kaggle.com/models/google/gemma",
         },
-        "kaggle_handle": "kaggle://keras/gemma/keras/gemma_instruct_2b_en/2",
+        "kaggle_handle": "kaggle://keras/gemma/keras/gemma_instruct_2b_en/3",
     },
     "gemma_1.1_instruct_2b_en": {
         "metadata": {
@@ -31,11 +27,9 @@ backbone_presets = {
                 "The 1.1 update improves model quality."
             ),
             "params": 2506172416,
-            "official_name": "Gemma",
             "path": "gemma",
-            "model_card": "https://www.kaggle.com/models/google/gemma",
         },
-        "kaggle_handle": "kaggle://keras/gemma/keras/gemma_1.1_instruct_2b_en/3",
+        "kaggle_handle": "kaggle://keras/gemma/keras/gemma_1.1_instruct_2b_en/4",
     },
     "code_gemma_1.1_2b_en": {
         "metadata": {
@@ -45,11 +39,9 @@ backbone_presets = {
                 "completion. The 1.1 update improves model quality."
             ),
             "params": 2506172416,
-            "official_name": "Gemma",
             "path": "gemma",
-            "model_card": "https://www.kaggle.com/models/google/gemma",
         },
-        "kaggle_handle": "kaggle://keras/codegemma/keras/code_gemma_1.1_2b_en/1",
+        "kaggle_handle": "kaggle://keras/codegemma/keras/code_gemma_1.1_2b_en/2",
     },
     "code_gemma_2b_en": {
         "metadata": {
@@ -59,21 +51,17 @@ backbone_presets = {
                 "completion."
             ),
             "params": 2506172416,
-            "official_name": "Gemma",
             "path": "gemma",
-            "model_card": "https://www.kaggle.com/models/google/gemma",
         },
-        "kaggle_handle": "kaggle://keras/codegemma/keras/code_gemma_2b_en/1",
+        "kaggle_handle": "kaggle://keras/codegemma/keras/code_gemma_2b_en/2",
     },
     "gemma_7b_en": {
         "metadata": {
             "description": "7 billion parameter, 28-layer, base Gemma model.",
             "params": 8537680896,
-            "official_name": "Gemma",
             "path": "gemma",
-            "model_card": "https://www.kaggle.com/models/google/gemma",
         },
-        "kaggle_handle": "kaggle://keras/gemma/keras/gemma_7b_en/2",
+        "kaggle_handle": "kaggle://keras/gemma/keras/gemma_7b_en/3",
     },
     "gemma_instruct_7b_en": {
         "metadata": {
@@ -81,11 +69,9 @@ backbone_presets = {
                 "7 billion parameter, 28-layer, instruction tuned Gemma model."
             ),
             "params": 8537680896,
-            "official_name": "Gemma",
             "path": "gemma",
-            "model_card": "https://www.kaggle.com/models/google/gemma",
         },
-        "kaggle_handle": "kaggle://keras/gemma/keras/gemma_instruct_7b_en/2",
+        "kaggle_handle": "kaggle://keras/gemma/keras/gemma_instruct_7b_en/3",
     },
     "gemma_1.1_instruct_7b_en": {
         "metadata": {
@@ -94,11 +80,9 @@ backbone_presets = {
                 "The 1.1 update improves model quality."
             ),
             "params": 8537680896,
-            "official_name": "Gemma",
             "path": "gemma",
-            "model_card": "https://www.kaggle.com/models/google/gemma",
         },
-        "kaggle_handle": "kaggle://keras/gemma/keras/gemma_1.1_instruct_7b_en/3",
+        "kaggle_handle": "kaggle://keras/gemma/keras/gemma_1.1_instruct_7b_en/4",
     },
     "code_gemma_7b_en": {
         "metadata": {
@@ -108,11 +92,9 @@ backbone_presets = {
                 "completion."
             ),
             "params": 8537680896,
-            "official_name": "Gemma",
             "path": "gemma",
-            "model_card": "https://www.kaggle.com/models/google/gemma",
         },
-        "kaggle_handle": "kaggle://keras/codegemma/keras/code_gemma_7b_en/1",
+        "kaggle_handle": "kaggle://keras/codegemma/keras/code_gemma_7b_en/2",
     },
     "code_gemma_instruct_7b_en": {
         "metadata": {
@@ -122,11 +104,9 @@ backbone_presets = {
                 "to code."
             ),
             "params": 8537680896,
-            "official_name": "Gemma",
             "path": "gemma",
-            "model_card": "https://www.kaggle.com/models/google/gemma",
         },
-        "kaggle_handle": "kaggle://keras/codegemma/keras/code_gemma_instruct_7b_en/1",
+        "kaggle_handle": "kaggle://keras/codegemma/keras/code_gemma_instruct_7b_en/2",
     },
     "code_gemma_1.1_instruct_7b_en": {
         "metadata": {
@@ -136,100 +116,86 @@ backbone_presets = {
                 "to code. The 1.1 update improves model quality."
             ),
             "params": 8537680896,
-            "official_name": "Gemma",
             "path": "gemma",
-            "model_card": "https://www.kaggle.com/models/google/gemma",
         },
-        "kaggle_handle": "kaggle://keras/codegemma/keras/code_gemma_1.1_instruct_7b_en/1",
+        "kaggle_handle": "kaggle://keras/codegemma/keras/code_gemma_1.1_instruct_7b_en/2",
     },
     "gemma2_2b_en": {
         "metadata": {
             "description": "2 billion parameter, 26-layer, base Gemma model.",
             "params": 2614341888,
-            "official_name": "Gemma",
             "path": "gemma",
-            "model_card": "https://www.kaggle.com/models/google/gemma",
         },
-        "kaggle_handle": "kaggle://keras/gemma2/keras/gemma2_2b_en/1",
+        "kaggle_handle": "kaggle://keras/gemma2/keras/gemma2_2b_en/2",
     },
     "gemma2_instruct_2b_en": {
         "metadata": {
-            "description": "2 billion parameter, 26-layer, instruction tuned Gemma model.",
+            "description": (
+                "2 billion parameter, 26-layer, instruction tuned Gemma model."
+            ),
             "params": 2614341888,
-            "official_name": "Gemma",
             "path": "gemma",
-            "model_card": "https://www.kaggle.com/models/google/gemma",
         },
-        "kaggle_handle": "kaggle://keras/gemma2/keras/gemma2_instruct_2b_en/1",
+        "kaggle_handle": "kaggle://keras/gemma2/keras/gemma2_instruct_2b_en/2",
     },
     "gemma2_9b_en": {
         "metadata": {
             "description": "9 billion parameter, 42-layer, base Gemma model.",
             "params": 9241705984,
-            "official_name": "Gemma",
             "path": "gemma",
-            "model_card": "https://www.kaggle.com/models/google/gemma",
         },
-        "kaggle_handle": "kaggle://keras/gemma2/keras/gemma2_9b_en/2",
+        "kaggle_handle": "kaggle://keras/gemma2/keras/gemma2_9b_en/3",
     },
     "gemma2_instruct_9b_en": {
         "metadata": {
-            "description": "9 billion parameter, 42-layer, instruction tuned Gemma model.",
+            "description": (
+                "9 billion parameter, 42-layer, instruction tuned Gemma model."
+            ),
             "params": 9241705984,
-            "official_name": "Gemma",
             "path": "gemma",
-            "model_card": "https://www.kaggle.com/models/google/gemma",
         },
-        "kaggle_handle": "kaggle://keras/gemma2/keras/gemma2_instruct_9b_en/2",
+        "kaggle_handle": "kaggle://keras/gemma2/keras/gemma2_instruct_9b_en/3",
     },
     "gemma2_27b_en": {
         "metadata": {
             "description": "27 billion parameter, 42-layer, base Gemma model.",
             "params": 27227128320,
-            "official_name": "Gemma",
             "path": "gemma",
-            "model_card": "https://www.kaggle.com/models/google/gemma",
         },
-        "kaggle_handle": "kaggle://keras/gemma2/keras/gemma2_27b_en/1",
+        "kaggle_handle": "kaggle://keras/gemma2/keras/gemma2_27b_en/2",
     },
     "gemma2_instruct_27b_en": {
         "metadata": {
-            "description": "27 billion parameter, 42-layer, instruction tuned Gemma model.",
+            "description": (
+                "27 billion parameter, 42-layer, instruction tuned Gemma model."
+            ),
             "params": 27227128320,
-            "official_name": "Gemma",
             "path": "gemma",
-            "model_card": "https://www.kaggle.com/models/google/gemma",
         },
-        "kaggle_handle": "kaggle://keras/gemma2/keras/gemma2_instruct_27b_en/1",
+        "kaggle_handle": "kaggle://keras/gemma2/keras/gemma2_instruct_27b_en/2",
     },
     "shieldgemma_2b_en": {
         "metadata": {
             "description": "2 billion parameter, 26-layer, ShieldGemma model.",
             "params": 2614341888,
-            "official_name": "Gemma",
             "path": "gemma",
-            "model_card": "https://www.kaggle.com/models/google/gemma",
         },
-        "kaggle_handle": "kaggle://google/shieldgemma/keras/shieldgemma_2b_en/1",
+        "kaggle_handle": "kaggle://google/shieldgemma/keras/shieldgemma_2b_en/2",
     },
     "shieldgemma_9b_en": {
         "metadata": {
             "description": "9 billion parameter, 42-layer, ShieldGemma model.",
             "params": 9241705984,
-            "official_name": "Gemma",
             "path": "gemma",
-            "model_card": "https://www.kaggle.com/models/google/gemma",
         },
-        "kaggle_handle": "kaggle://google/shieldgemma/keras/shieldgemma_9b_en/1",
+        "kaggle_handle": "kaggle://google/shieldgemma/keras/shieldgemma_9b_en/2",
     },
     "shieldgemma_27b_en": {
         "metadata": {
             "description": "27 billion parameter, 42-layer, ShieldGemma model.",
             "params": 27227128320,
-            "official_name": "Gemma",
             "path": "gemma",
-            "model_card": "https://www.kaggle.com/models/google/gemma",
         },
-        "kaggle_handle": "kaggle://google/shieldgemma/keras/shieldgemma_27b_en/1",
+        "kaggle_handle": "kaggle://google/shieldgemma/keras/shieldgemma_27b_en/2",
     },
 }

keras_hub/src/models/gpt2/gpt2_causal_lm.py

@@ -172,8 +172,8 @@ class GPT2CausalLM(CausalLM):
         Args:
             token_ids: a dense int Tensor with shape `(batch_size, max_length)`.
             cache: a dense float Tensor, the cache of key and value.
-            cache_update_index: int, or int Tensor. The index of current inputs in the
-                whole sequence.
+            cache_update_index: int, or int Tensor. The index of current inputs
+                in the whole sequence.
 
         Returns:
             A (logits, hidden_states, cache) tuple. Where `logits` is the

keras_hub/src/models/gpt2/gpt2_presets.py

@@ -9,11 +9,9 @@ backbone_presets = {
                 "Trained on WebText."
             ),
             "params": 124439808,
-            "official_name": "GPT-2",
             "path": "gpt2",
-            "model_card": "https://github.com/openai/gpt-2/blob/master/model_card.md",
         },
-        "kaggle_handle": "kaggle://keras/gpt2/keras/gpt2_base_en/2",
+        "kaggle_handle": "kaggle://keras/gpt2/keras/gpt2_base_en/3",
     },
     "gpt2_medium_en": {
         "metadata": {
@@ -22,11 +20,9 @@ backbone_presets = {
                 "Trained on WebText."
             ),
             "params": 354823168,
-            "official_name": "GPT-2",
             "path": "gpt2",
-            "model_card": "https://github.com/openai/gpt-2/blob/master/model_card.md",
         },
-        "kaggle_handle": "kaggle://keras/gpt2/keras/gpt2_medium_en/2",
+        "kaggle_handle": "kaggle://keras/gpt2/keras/gpt2_medium_en/3",
     },
     "gpt2_large_en": {
         "metadata": {
@@ -35,11 +31,9 @@ backbone_presets = {
                 "Trained on WebText."
             ),
             "params": 774030080,
-            "official_name": "GPT-2",
             "path": "gpt2",
-            "model_card": "https://github.com/openai/gpt-2/blob/master/model_card.md",
         },
-        "kaggle_handle": "kaggle://keras/gpt2/keras/gpt2_large_en/2",
+        "kaggle_handle": "kaggle://keras/gpt2/keras/gpt2_large_en/3",
     },
     "gpt2_extra_large_en": {
         "metadata": {
@@ -48,11 +42,9 @@ backbone_presets = {
                 "Trained on WebText."
             ),
             "params": 1557611200,
-            "official_name": "GPT-2",
             "path": "gpt2",
-            "model_card": "https://github.com/openai/gpt-2/blob/master/model_card.md",
         },
-        "kaggle_handle": "kaggle://keras/gpt2/keras/gpt2_extra_large_en/2",
+        "kaggle_handle": "kaggle://keras/gpt2/keras/gpt2_extra_large_en/3",
     },
     "gpt2_base_en_cnn_dailymail": {
         "metadata": {
@@ -61,9 +53,8 @@ backbone_presets = {
                 "Finetuned on the CNN/DailyMail summarization dataset."
             ),
             "params": 124439808,
-            "official_name": "GPT-2",
             "path": "gpt2",
         },
-        "kaggle_handle": "kaggle://keras/gpt2/keras/gpt2_base_en_cnn_dailymail/2",
+        "kaggle_handle": "kaggle://keras/gpt2/keras/gpt2_base_en_cnn_dailymail/3",
     },
 }

keras_hub/src/models/gpt_neo_x/gpt_neo_x_attention.py

@@ -202,7 +202,8 @@ class GPTNeoXAttention(keras.layers.Layer):
             training=training,
         )
 
-        # Reshape `attention_output` to `(batch_size, sequence_length, hidden_dim)`.
+        # Reshape `attention_output` to
+        # `(batch_size, sequence_length, hidden_dim)`.
         attention_output = ops.reshape(
             attention_output,
             [

keras_hub/src/models/gpt_neo_x/gpt_neo_x_causal_lm.py

@@ -27,9 +27,9 @@ class GPTNeoXCausalLM(CausalLM):
 
     Args:
         backbone: A `keras_hub.models.GPTNeoXBackbone` instance.
-        preprocessor: A `keras_hub.models.GPTNeoXCausalLMPreprocessor` or `None`.
-            If `None`, this model will not apply preprocessing, and inputs
-            should be preprocessed before calling the model.
+        preprocessor: A `keras_hub.models.GPTNeoXCausalLMPreprocessor` or
+            `None`. If `None`, this model will not apply preprocessing, and
+            inputs should be preprocessed before calling the model.
     """
 
     backbone_cls = GPTNeoXBackbone

keras_hub/src/models/gpt_neo_x/gpt_neo_x_decoder.py

@@ -16,7 +16,8 @@ class GPTNeoXDecoder(keras.layers.Layer):
 
     This class follows the architecture of the GPT-NeoX decoder layer in the
     paper [GPT-NeoX-20B: An Open-Source Autoregressive Language Model](https://arxiv.org/abs/2204.06745).
-    Users can instantiate multiple instances of this class to stack up a decoder.
+    Users can instantiate multiple instances of this class to stack up a
+    decoder.
 
     This layer will always apply a causal mask to the decoder attention layer.
 

keras_hub/src/models/image_classifier.py

@@ -15,11 +15,156 @@ class ImageClassifier(Task):
 
     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.
 
-    All `ImageClassifier` tasks include a `from_preset()` constructor which can be
-    used to load a pre-trained config and weights.
+    Args:
+        backbone: A `keras_hub.models.Backbone` 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: `"avg"` or `"max"`. The type of pooling to apply on backbone
+            output. Defaults to average pooling.
+        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.ImageClassifier.from_preset(
+        "resnet_50_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.ImageClassifier.from_preset(
+        "resnet_50_imagenet"
+    )
+    classifier.fit(x=images, y=labels, batch_size=2)
+    ```
+
+    Call `fit()` with custom loss, optimizer and backbone.
+    ```python
+    classifier = keras_hub.models.ImageClassifier.from_preset(
+        "resnet_50_imagenet"
+    )
+    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]
+    backbone = keras_hub.models.ResNetBackbone(
+        stackwise_num_filters=[64, 64, 64],
+        stackwise_num_blocks=[2, 2, 2],
+        stackwise_num_strides=[1, 2, 2],
+        block_type="basic_block",
+        use_pre_activation=True,
+        pooling="avg",
+    )
+    classifier = keras_hub.models.ImageClassifier(
+        backbone=backbone,
+        num_classes=4,
+    )
+    classifier.fit(x=images, y=labels, batch_size=2)
+    ```
     """
 
+    def __init__(
+        self,
+        backbone,
+        num_classes,
+        preprocessor=None,
+        pooling="avg",
+        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.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",
+            )
+        else:
+            raise ValueError(
+                "Unknown `pooling` type. Polling should be either `'avg'` or "
+                f"`'max'`. Received: pooling={pooling}."
+            )
+        self.output_dropout = keras.layers.Dropout(
+            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)
+        x = self.pooler(x)
+        x = self.output_dropout(x)
+        outputs = self.output_dense(x)
+        super().__init__(
+            inputs=inputs,
+            outputs=outputs,
+            **kwargs,
+        )
+
+        # === Config ===
+        self.num_classes = num_classes
+        self.activation = activation
+        self.pooling = pooling
+        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,
+                "activation": self.activation,
+                "dropout": self.dropout,
+            }
+        )
+        return config
+
     def compile(
         self,
         optimizer="auto",

keras_hub/src/models/image_classifier_preprocessor.py

@@ -38,15 +38,18 @@ class ImageClassifierPreprocessor(Preprocessor):
     )
 
     # Resize a single image for resnet 50.
-    x = np.ones((512, 512, 3))
+    x = np.random.randint(0, 256, (512, 512, 3))
     x = preprocessor(x)
 
     # Resize a labeled image.
-    x, y = np.ones((512, 512, 3)), 1
+    x, y = np.random.randint(0, 256, (512, 512, 3)), 1
     x, y = preprocessor(x, y)
 
     # Resize a batch of labeled images.
-    x, y = [np.ones((512, 512, 3)), np.zeros((512, 512, 3))], [1, 0]
+    x, y = [
+        np.random.randint(0, 256, (512, 512, 3)),
+        np.zeros((512, 512, 3))
+    ], [1, 0]
     x, y = preprocessor(x, y)
 
     # Use a `tf.data.Dataset`.