Add pyramid output for densenet, cspDarknet

keras_nlp/src/models/csp_darknet/csp_darknet_backbone.py

@@ -15,11 +15,11 @@
 from keras import layers
 
 from keras_nlp.src.api_export import keras_nlp_export
-from keras_nlp.src.models.backbone import Backbone
+from keras_nlp.src.models.feature_pyramid_backbone import FeaturePyramidBackbone
 
 
 @keras_nlp_export("keras_nlp.models.CSPDarkNetBackbone")
-class CSPDarkNetBackbone(Backbone):
+class CSPDarkNetBackbone(FeaturePyramidBackbone):
     """This class represents Keras Backbone of CSPDarkNet model.
 
     This class implements a CSPDarkNet backbone as described in
@@ -39,7 +39,7 @@ class CSPDarkNetBackbone(Backbone):
             `"basic_block"` for basic conv block.
             Defaults to "basic_block".
         image_shape: tuple. The input shape without the batch size.
-            Defaults to `(None, None, 3)`.
+            Defaults to `(224, 224, 3)`.
 
     Examples:
     ```python
@@ -87,6 +87,8 @@ def __init__(
         x = apply_darknet_conv_block(
             base_channels, kernel_size=3, strides=1, name="stem_conv"
         )(x)
+
+        pyramid_outputs = {}
         for index, (channels, depth) in enumerate(
             zip(stackwise_num_filters, stackwise_depth)
         ):
@@ -111,6 +113,7 @@ def __init__(
                 residual=(index != len(stackwise_depth) - 1),
                 name=f"dark{index + 2}_csp",
             )(x)
+            pyramid_outputs[f"P{index + 2}"] = x
 
         super().__init__(inputs=image_input, outputs=x, **kwargs)
 
@@ -120,6 +123,7 @@ def __init__(
         self.include_rescaling = include_rescaling
         self.block_type = block_type
         self.image_shape = image_shape
+        self.pyramid_outputs = pyramid_outputs
 
     def get_config(self):
         config = super().get_config()

keras_nlp/src/models/csp_darknet/csp_darknet_backbone_test.py

@@ -24,23 +24,34 @@
 class CSPDarkNetBackboneTest(TestCase):
     def setUp(self):
         self.init_kwargs = {
-            "stackwise_num_filters": [32, 64, 128, 256],
+            "stackwise_num_filters": [2, 4, 6, 8],
             "stackwise_depth": [1, 3, 3, 1],
             "include_rescaling": False,
             "block_type": "basic_block",
-            "image_shape": (224, 224, 3),
+            "image_shape": (32, 32, 3),
         }
-        self.input_data = np.ones((2, 224, 224, 3), dtype="float32")
+        self.input_size = 32
+        self.input_data = np.ones(
+            (2, self.input_size, self.input_size, 3), dtype="float32"
+        )
 
     def test_backbone_basics(self):
         self.run_backbone_test(
             cls=CSPDarkNetBackbone,
             init_kwargs=self.init_kwargs,
             input_data=self.input_data,
-            expected_output_shape=(2, 7, 7, 256),
+            expected_output_shape=(2, 1, 1, 8),
             run_mixed_precision_check=False,
         )
 
+    def test_pyramid_output_format(self):
+        self.run_pyramid_output_test(
+            cls=CSPDarkNetBackbone,
+            init_kwargs=self.init_kwargs,
+            input_data=self.input_data,
+            expected_pyramid_output_keys=["P2", "P3", "P4", "P5"],
+        )
+
     @pytest.mark.large
     def test_saved_model(self):
         self.run_model_saving_test(

keras_nlp/src/models/densenet/densenet_backbone.py

@@ -14,14 +14,14 @@
 import keras
 
 from keras_nlp.src.api_export import keras_nlp_export
-from keras_nlp.src.models.backbone import Backbone
+from keras_nlp.src.models.feature_pyramid_backbone import FeaturePyramidBackbone
 
 BN_AXIS = 3
 BN_EPSILON = 1.001e-5
 
 
 @keras_nlp_export("keras_nlp.models.DenseNetBackbone")
-class DenseNetBackbone(Backbone):
+class DenseNetBackbone(FeaturePyramidBackbone):
     """Instantiates the DenseNet architecture.
 
     This class implements a DenseNet backbone as described in
@@ -85,6 +85,7 @@ def __init__(
             3, strides=2, padding="same", name="pool1"
         )(x)
 
+        pyramid_outputs = {}
         for stack_index in range(len(stackwise_num_repeats) - 1):
             index = stack_index + 2
             x = apply_dense_block(
@@ -93,6 +94,7 @@ def __init__(
                 growth_rate,
                 name=f"conv{index}",
             )
+            pyramid_outputs[f"P{index}"] = x
             x = apply_transition_block(
                 x, compression_ratio, name=f"pool{index}"
             )
@@ -103,7 +105,7 @@ def __init__(
             growth_rate,
             name=f"conv{len(stackwise_num_repeats) + 1}",
         )
-
+        pyramid_outputs[f"P{len(stackwise_num_repeats) +1}"] = x
         x = keras.layers.BatchNormalization(
             axis=BN_AXIS, epsilon=BN_EPSILON, name="bn"
         )(x)
@@ -117,6 +119,7 @@ def __init__(
         self.compression_ratio = compression_ratio
         self.growth_rate = growth_rate
         self.image_shape = image_shape
+        self.pyramid_outputs = pyramid_outputs
 
     def get_config(self):
         config = super().get_config()

keras_nlp/src/models/densenet/densenet_backbone_test.py

@@ -22,23 +22,34 @@
 class DenseNetBackboneTest(TestCase):
     def setUp(self):
         self.init_kwargs = {
-            "stackwise_num_repeats": [6, 12, 24, 16],
+            "stackwise_num_repeats": [2, 4, 6, 4],
             "include_rescaling": True,
             "compression_ratio": 0.5,
-            "growth_rate": 32,
-            "image_shape": (224, 224, 3),
+            "growth_rate": 2,
+            "image_shape": (32, 32, 3),
         }
-        self.input_data = np.ones((2, 224, 224, 3), dtype="float32")
+        self.input_size = 32
+        self.input_data = np.ones(
+            (2, self.input_size, self.input_size, 3), dtype="float32"
+        )
 
     def test_backbone_basics(self):
         self.run_backbone_test(
             cls=DenseNetBackbone,
             init_kwargs=self.init_kwargs,
             input_data=self.input_data,
-            expected_output_shape=(2, 7, 7, 1024),
+            expected_output_shape=(2, 1, 1, 24),
             run_mixed_precision_check=False,
         )
 
+    def test_pyramid_output_format(self):
+        self.run_pyramid_output_test(
+            cls=DenseNetBackbone,
+            init_kwargs=self.init_kwargs,
+            input_data=self.input_data,
+            expected_pyramid_output_keys=["P2", "P3", "P4", "P5"],
+        )
+
     @pytest.mark.large
     def test_saved_model(self):
         self.run_model_saving_test(

keras_nlp/src/models/resnet/resnet_backbone_test.py

@@ -14,7 +14,6 @@
 
 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
@@ -58,18 +57,12 @@ def test_pyramid_output_format(self):
         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.run_pyramid_output_test(
+            cls=ResNetBackbone,
+            init_kwargs=init_kwargs,
+            input_data=self.input_data,
+            expected_pyramid_output_keys=["P2", "P3", "P4"],
         )
-        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"),

keras_nlp/src/tests/test_case.py

@@ -604,5 +604,27 @@ def compare(actual, expected):
 
             tree.map_structure(compare, output, expected_partial_output)
 
+    def run_pyramid_output_test(
+        self,
+        cls,
+        init_kwargs,
+        input_data,
+        expected_pyramid_output_keys,
+    ):
+        """Run Tests for Feature Pyramid output keys and shape."""
+        backbone = cls(**init_kwargs)
+        model = keras.models.Model(backbone.inputs, backbone.pyramid_outputs)
+        output_data = model(input_data)
+
+        self.assertIsInstance(output_data, dict)
+        self.assertEqual(
+            list(output_data.keys()), list(backbone.pyramid_outputs.keys())
+        )
+        self.assertEqual(list(output_data.keys()), expected_pyramid_output_keys)
+
+        for k, v in output_data.items():
+            size = input_data.shape[1] // (2 ** int(k[1:]))
+            self.assertEqual(tuple(v.shape[:3]), (2, size, size))
+
     def get_test_data_dir(self):
         return str(pathlib.Path(__file__).parent / "test_data")