keras-hub-nightly 0.16.1.dev202410020340__py3-none-any.whl → 0.16.1.dev202410040340__py3-none-any.whl

keras_hub/src/models/vgg/vgg_backbone.py

@@ -21,17 +21,6 @@ class VGGBackbone(Backbone):
             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.
 
     Examples:
     ```python
@@ -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,8 +44,7 @@ class VGGBackbone(Backbone):
         self,
         stackwise_num_repeats,
         stackwise_num_filters,
-        image_shape=(224, 224, 3),
-        pooling="avg",
+        image_shape=(None, None, 3),
         **kwargs,
     ):
 
@@ -76,10 +63,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 +70,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,90 @@ 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
 
 
 @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,
@@ -69,26 +101,67 @@ class VGGImageClassifier(ImageClassifier):
         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="flatten",
+        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.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.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)
-
-        # Instantiate using Functional API Model constructor
-        super().__init__(
+        # Skip the parent class functional model.
+        Task.__init__(
+            self,
             inputs=inputs,
             outputs=outputs,
             **kwargs,
@@ -97,6 +170,9 @@ 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
 
     def get_config(self):
         # Backbone serialized in `super`
@@ -104,7 +180,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/tokenizers/tokenizer.py

@@ -10,7 +10,7 @@ from keras_hub.src.utils.preset_utils import builtin_presets
 from keras_hub.src.utils.preset_utils import find_subclass
 from keras_hub.src.utils.preset_utils import get_file
 from keras_hub.src.utils.preset_utils import get_preset_loader
-from keras_hub.src.utils.preset_utils import save_serialized_object
+from keras_hub.src.utils.preset_utils import get_preset_saver
 from keras_hub.src.utils.python_utils import classproperty
 from keras_hub.src.utils.tensor_utils import preprocessing_function
 
@@ -189,11 +189,8 @@ class Tokenizer(PreprocessingLayer):
         Args:
             preset_dir: The path to the local model preset directory.
         """
-        save_serialized_object(self, preset_dir, config_file=self.config_name)
-        subdir = self.config_name.split(".")[0]
-        asset_dir = os.path.join(preset_dir, ASSET_DIR, subdir)
-        os.makedirs(asset_dir, exist_ok=True)
-        self.save_assets(asset_dir)
+        saver = get_preset_saver(preset_dir)
+        saver.save_tokenizer(self)
 
     @preprocessing_function
     def call(self, inputs, *args, training=None, **kwargs):

keras_hub/src/utils/preset_utils.py

@@ -267,64 +267,6 @@ def check_file_exists(preset, path):
     return True
 
 
-def get_tokenizer(layer):
-    """Get the tokenizer from any KerasHub model or layer."""
-    # Avoid circular import.
-    from keras_hub.src.tokenizers.tokenizer import Tokenizer
-
-    if isinstance(layer, Tokenizer):
-        return layer
-    if hasattr(layer, "tokenizer"):
-        return layer.tokenizer
-    if hasattr(layer, "preprocessor"):
-        return getattr(layer.preprocessor, "tokenizer", None)
-    return None
-
-
-def recursive_pop(config, key):
-    """Remove a key from a nested config object"""
-    config.pop(key, None)
-    for value in config.values():
-        if isinstance(value, dict):
-            recursive_pop(value, key)
-
-
-# TODO: refactor saving routines into a PresetSaver class?
-def make_preset_dir(preset):
-    os.makedirs(preset, exist_ok=True)
-
-
-def save_serialized_object(
-    layer,
-    preset,
-    config_file=CONFIG_FILE,
-    config_to_skip=[],
-):
-    make_preset_dir(preset)
-    config_path = os.path.join(preset, config_file)
-    config = keras.saving.serialize_keras_object(layer)
-    config_to_skip += ["compile_config", "build_config"]
-    for c in config_to_skip:
-        recursive_pop(config, c)
-    with open(config_path, "w") as config_file:
-        config_file.write(json.dumps(config, indent=4))
-
-
-def save_metadata(layer, preset):
-    from keras_hub.src.version_utils import __version__ as keras_hub_version
-
-    keras_version = keras.version() if hasattr(keras, "version") else None
-    metadata = {
-        "keras_version": keras_version,
-        "keras_hub_version": keras_hub_version,
-        "parameter_count": layer.count_params(),
-        "date_saved": datetime.datetime.now().strftime("%Y-%m-%d@%H:%M:%S"),
-    }
-    metadata_path = os.path.join(preset, METADATA_FILE)
-    with open(metadata_path, "w") as metadata_file:
-        metadata_file.write(json.dumps(metadata, indent=4))
-
-
 def _validate_backbone(preset):
     config_path = os.path.join(preset, CONFIG_FILE)
     if not os.path.exists(config_path):
@@ -518,6 +460,8 @@ def load_serialized_object(config, **kwargs):
 def check_config_class(config):
     """Validate a preset is being loaded on the correct class."""
     registered_name = config["registered_name"]
+    if registered_name in ("Functional", "Sequential"):
+        return keras.Model
     cls = keras.saving.get_registered_object(registered_name)
     if cls is None:
         raise ValueError(
@@ -600,6 +544,13 @@ def get_preset_loader(preset):
         )
 
 
+def get_preset_saver(preset):
+    # Unlike loading, we only support one form of saving; Keras serialized
+    # configs and saved weights. We keep the rough API structure as loading
+    # just for simplicity.
+    return KerasPresetSaver(preset)
+
+
 class PresetLoader:
     def __init__(self, preset, config):
         self.config = config
@@ -684,7 +635,8 @@ class KerasPresetLoader(PresetLoader):
     def load_tokenizer(self, cls, config_name=TOKENIZER_CONFIG_FILE, **kwargs):
         tokenizer_config = load_json(self.preset, config_name)
         tokenizer = load_serialized_object(tokenizer_config, **kwargs)
-        tokenizer.load_preset_assets(self.preset)
+        if hasattr(tokenizer, "load_preset_assets"):
+            tokenizer.load_preset_assets(self.preset)
         return tokenizer
 
     def load_audio_converter(self, cls, **kwargs):
@@ -709,7 +661,9 @@ class KerasPresetLoader(PresetLoader):
             )
         # We found a `task.json` with a complete config for our class.
         task = load_serialized_object(task_config, **kwargs)
-        if task.preprocessor:
+        if task.preprocessor and hasattr(
+            task.preprocessor, "load_preset_assets"
+        ):
             task.preprocessor.load_preset_assets(self.preset)
         if load_weights:
             has_task_weights = check_file_exists(self.preset, TASK_WEIGHTS_FILE)
@@ -735,5 +689,93 @@ class KerasPresetLoader(PresetLoader):
             return super().load_preprocessor(cls, **kwargs)
         # We found a `preprocessing.json` with a complete config for our class.
         preprocessor = load_serialized_object(preprocessor_json, **kwargs)
-        preprocessor.load_preset_assets(self.preset)
+        if hasattr(preprocessor, "load_preset_assets"):
+            preprocessor.load_preset_assets(self.preset)
         return preprocessor
+
+
+class KerasPresetSaver:
+    def __init__(self, preset_dir):
+        os.makedirs(preset_dir, exist_ok=True)
+        self.preset_dir = preset_dir
+
+    def save_backbone(self, backbone):
+        self._save_serialized_object(backbone, config_file=CONFIG_FILE)
+        backbone_weight_path = os.path.join(self.preset_dir, MODEL_WEIGHTS_FILE)
+        backbone.save_weights(backbone_weight_path)
+        self._save_metadata(backbone)
+
+    def save_tokenizer(self, tokenizer):
+        config_file = TOKENIZER_CONFIG_FILE
+        if hasattr(tokenizer, "config_file"):
+            config_file = tokenizer.config_file
+        self._save_serialized_object(tokenizer, config_file)
+        # Save assets.
+        subdir = config_file.split(".")[0]
+        asset_dir = os.path.join(self.preset_dir, ASSET_DIR, subdir)
+        os.makedirs(asset_dir, exist_ok=True)
+        tokenizer.save_assets(asset_dir)
+
+    def save_audio_converter(self, converter):
+        self._save_serialized_object(converter, AUDIO_CONVERTER_CONFIG_FILE)
+
+    def save_image_converter(self, converter):
+        self._save_serialized_object(converter, IMAGE_CONVERTER_CONFIG_FILE)
+
+    def save_task(self, task):
+        # Save task specific config and weights.
+        self._save_serialized_object(task, TASK_CONFIG_FILE)
+        if task.has_task_weights():
+            task_weight_path = os.path.join(self.preset_dir, TASK_WEIGHTS_FILE)
+            task.save_task_weights(task_weight_path)
+        # Save backbone.
+        if hasattr(task.backbone, "save_to_preset"):
+            task.backbone.save_to_preset(self.preset_dir)
+        else:
+            # Allow saving a `keras.Model` that is not a backbone subclass.
+            self.save_backbone(task.backbone)
+        # Save preprocessor.
+        if task.preprocessor and hasattr(task.preprocessor, "save_to_preset"):
+            task.preprocessor.save_to_preset(self.preset_dir)
+        else:
+            # Allow saving a `keras.Layer` that is not a preprocessor subclass.
+            self.save_preprocessor(task.preprocessor)
+
+    def save_preprocessor(self, preprocessor):
+        config_file = PREPROCESSOR_CONFIG_FILE
+        if hasattr(preprocessor, "config_file"):
+            config_file = preprocessor.config_file
+        self._save_serialized_object(preprocessor, config_file)
+        for layer in preprocessor._flatten_layers(include_self=False):
+            if hasattr(layer, "save_to_preset"):
+                layer.save_to_preset(self.preset_dir)
+
+    def _recursive_pop(self, config, key):
+        """Remove a key from a nested config object"""
+        config.pop(key, None)
+        for value in config.values():
+            if isinstance(value, dict):
+                self._recursive_pop(value, key)
+
+    def _save_serialized_object(self, layer, config_file):
+        config_path = os.path.join(self.preset_dir, config_file)
+        config = keras.saving.serialize_keras_object(layer)
+        config_to_skip = ["compile_config", "build_config"]
+        for key in config_to_skip:
+            self._recursive_pop(config, key)
+        with open(config_path, "w") as config_file:
+            config_file.write(json.dumps(config, indent=4))
+
+    def _save_metadata(self, layer):
+        from keras_hub.src.version_utils import __version__ as keras_hub_version
+
+        keras_version = keras.version() if hasattr(keras, "version") else None
+        metadata = {
+            "keras_version": keras_version,
+            "keras_hub_version": keras_hub_version,
+            "parameter_count": layer.count_params(),
+            "date_saved": datetime.datetime.now().strftime("%Y-%m-%d@%H:%M:%S"),
+        }
+        metadata_path = os.path.join(self.preset_dir, METADATA_FILE)
+        with open(metadata_path, "w") as metadata_file:
+            metadata_file.write(json.dumps(metadata, indent=4))

keras_hub/src/utils/timm/preset_loader.py

@@ -14,7 +14,7 @@ class TimmPresetLoader(PresetLoader):
         architecture = self.config["architecture"]
         if "resnet" in architecture:
             self.converter = convert_resnet
-        if "densenet" in architecture:
+        elif "densenet" in architecture:
             self.converter = convert_densenet
         else:
             raise ValueError(
@@ -52,20 +52,19 @@ class TimmPresetLoader(PresetLoader):
         pretrained_cfg = self.config.get("pretrained_cfg", None)
         if not pretrained_cfg or "input_size" not in pretrained_cfg:
             return None
-        # This assumes the same basic setup for all timm preprocessing, and that
-        # all our image conversion will be via a `ResizingImageConverter. We may
+        # This assumes the same basic setup for all timm preprocessing, We may
         # need to extend this as we cover more model types.
         input_size = pretrained_cfg["input_size"]
         mean = pretrained_cfg["mean"]
-        variance = [s**2 for s in pretrained_cfg["std"]]
+        std = pretrained_cfg["std"]
+        scale = [1.0 / 255.0 / s for s in std]
+        offset = [-m / s for m, s in zip(mean, std)]
         interpolation = pretrained_cfg["interpolation"]
         if interpolation not in ("bilinear", "nearest", "bicubic"):
             interpolation = "bilinear"  # Unsupported interpolation type.
         return cls(
-            width=input_size[1],
-            height=input_size[2],
-            scale=1 / 255.0,
-            mean=mean,
-            variance=variance,
+            image_size=input_size[1:],
+            scale=scale,
+            offset=offset,
             interpolation=interpolation,
         )

keras_hub/src/version_utils.py

@@ -1,7 +1,7 @@
 from keras_hub.src.api_export import keras_hub_export
 
 # Unique source of truth for the version number.
-__version__ = "0.16.1.dev202410020340"
+__version__ = "0.16.1.dev202410040340"
 
 
 @keras_hub_export("keras_hub.version")

{keras_hub_nightly-0.16.1.dev202410020340.dist-info → keras_hub_nightly-0.16.1.dev202410040340.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: keras-hub-nightly
-Version: 0.16.1.dev202410020340
+Version: 0.16.1.dev202410040340
 Summary: Industry-strength Natural Language Processing extensions for Keras.
 Home-page: https://github.com/keras-team/keras-hub
 Author: Keras team