Address comments

Author: james77777778

keras_nlp/api/models/__init__.py

@@ -183,9 +183,6 @@
 from keras_nlp.src.models.phi3.phi3_tokenizer import Phi3Tokenizer
 from keras_nlp.src.models.preprocessor import Preprocessor
 from keras_nlp.src.models.resnet.resnet_backbone import ResNetBackbone
-from keras_nlp.src.models.resnet.resnet_feature_pyramid_backbone import (
-    ResNetFeaturePyramidBackbone,
-)
 from keras_nlp.src.models.resnet.resnet_image_classifier import (
     ResNetImageClassifier,
 )

keras_nlp/src/models/feature_pyramid_backbone.py

@@ -19,6 +19,26 @@
 
 @keras_nlp_export("keras_nlp.models.FeaturePyramidBackbone")
 class FeaturePyramidBackbone(Backbone):
+    """A backbone with feature pyramid outputs.
+
+    `FeaturePyramidBackbone` extends `Backbone` with a single `pyramid_outputs`
+    property for accessing the feature pyramid outputs of the model. Subclassers
+    should set the `pyramid_outputs` property during the model constructor.
+
+    Example:
+
+    ```python
+    input_data = np.random.uniform(0, 255, size=(2, 224, 224, 3))
+
+    # Convert to feature pyramid output format using ResNet.
+    backbone = ResNetBackbone.from_preset("resnet50")
+    model = keras.Model(
+        inputs=backbone.inputs, outputs=backbone.pyramid_outputs
+    )
+    model(input_data)  # A dict containing the keys ["P2", "P3", "P4", "P5"]
+    ```
+    """
+
     @property
     def pyramid_outputs(self):
         """A dict for feature pyramid outputs.

keras_nlp/src/models/resnet/resnet_backbone_test.py

@@ -14,6 +14,7 @@
 
 import pytest
 from absl.testing import parameterized
+from keras import models
 from keras import ops
 
 from keras_nlp.src.models.resnet.resnet_backbone import ResNetBackbone
@@ -29,8 +30,8 @@ def setUp(self):
             "input_image_shape": (None, None, 3),
             "pooling": "avg",
         }
-        self.input_size = (16, 16)
-        self.input_data = ops.ones((2, 16, 16, 3))
+        self.input_size = 64
+        self.input_data = ops.ones((2, self.input_size, self.input_size, 3))
 
     @parameterized.named_parameters(
         ("v1_basic", False, "basic_block"),
@@ -52,6 +53,24 @@ def test_backbone_basics(self, use_pre_activation, block_type):
             ),
         )
 
+    def test_pyramid_output_format(self):
+        init_kwargs = self.init_kwargs.copy()
+        init_kwargs.update(
+            {"block_type": "basic_block", "use_pre_activation": False}
+        )
+        backbone = ResNetBackbone(**init_kwargs)
+        model = models.Model(backbone.inputs, backbone.pyramid_outputs)
+        output_data = model(self.input_data)
+
+        self.assertIsInstance(output_data, dict)
+        self.assertEqual(
+            list(output_data.keys()), list(backbone.pyramid_outputs.keys())
+        )
+        self.assertEqual(list(output_data.keys()), ["P2", "P3", "P4"])
+        for k, v in output_data.items():
+            size = self.input_size // (2 ** int(k[1:]))
+            self.assertEqual(tuple(v.shape[:3]), (2, size, size))
+
     @parameterized.named_parameters(
         ("v1_basic", False, "basic_block"),
         ("v1_bottleneck", False, "bottleneck_block"),
@@ -65,7 +84,7 @@ def test_saved_model(self, use_pre_activation, block_type):
             {
                 "block_type": block_type,
                 "use_pre_activation": use_pre_activation,
-                "input_image_shape": (16, 16, 3),
+                "input_image_shape": (None, None, 3),
             }
         )
         self.run_model_saving_test(

keras_nlp/src/models/resnet/resnet_feature_pyramid_backbone.py

@@ -29,7 +29,7 @@ class ResNetImageClassifier(ImageClassifier):
             the `Dense` layer. Set `activation=None` to return the output
             logits. Defaults to `"softmax"`.
         head_dtype: `None` or str or `keras.mixed_precision.DTypePolicy`. The
-            dtype to use for the head's computations and weights.
+            dtype to use for the classification head's computations and weights.
 
     To fine-tune with `fit()`, pass a dataset containing tuples of `(x, y)`
     where `x` is a tensor and `y` is a integer from `[0, num_classes)`.

keras_nlp/src/models/resnet/resnet_feature_pyramid_backbone_test.py

@@ -50,7 +50,6 @@
 KAGGLE_PREFIX = "kaggle://"
 GS_PREFIX = "gs://"
 HF_PREFIX = "hf://"
-TIMM_PREFIX = "hf://timm"
 
 KAGGLE_SCHEME = "kaggle"
 GS_SCHEME = "gs"
@@ -545,7 +544,9 @@ def check_format(preset):
     if check_file_exists(preset, SAFETENSOR_FILE) or check_file_exists(
         preset, SAFETENSOR_CONFIG_FILE
     ):
-        if TIMM_PREFIX in preset:
+        # Determine the format by parsing the config file.
+        config = load_config(preset, HF_CONFIG_FILE)
+        if "hf://timm" in preset or "architecture" in config:
             return "timm"
         return "transformers"
 

keras_nlp/src/models/resnet/resnet_image_classifier.py

@@ -16,7 +16,7 @@
 from keras_nlp.src.utils.preset_utils import HF_CONFIG_FILE
 from keras_nlp.src.utils.preset_utils import jax_memory_cleanup
 from keras_nlp.src.utils.preset_utils import load_config
-from keras_nlp.src.utils.timm.safetensor_utils import SafetensorLoader
+from keras_nlp.src.utils.transformers.safetensor_utils import SafetensorLoader
 
 
 def convert_backbone_config(timm_config):
@@ -60,14 +60,11 @@ def convert_backbone_config(timm_config):
 
 
 def convert_weights(backbone, loader, timm_config):
-    def transpose_conv2d(x, shape):
-        return np.transpose(x, (2, 3, 1, 0))
-
     def port_conv2d(keras_layer_name, hf_weight_prefix):
         loader.port_weight(
             backbone.get_layer(keras_layer_name).kernel,
             hf_weight_key=f"{hf_weight_prefix}.weight",
-            hook_fn=transpose_conv2d,
+            hook_fn=lambda x, _: np.transpose(x, (2, 3, 1, 0)),
         )
 
     def port_batch_normalization(keras_layer_name, hf_weight_prefix):
@@ -168,6 +165,7 @@ def load_resnet_backbone(cls, preset, load_weights, **kwargs):
     backbone = cls(**keras_config, **kwargs)
     if load_weights:
         jax_memory_cleanup(backbone)
-        with SafetensorLoader(preset) as loader:
+        # Use prefix="" to avoid using `get_prefixed_key`.
+        with SafetensorLoader(preset, prefix="") as loader:
             convert_weights(backbone, loader, timm_config)
     return backbone