From 351000d714efefe8480f11b04291003115f762dc Mon Sep 17 00:00:00 2001 From: Sachin Prasad Date: Tue, 27 Aug 2024 18:59:14 +0000 Subject: [PATCH 1/7] add pyramid outputs --- .../src/models/densenet/densenet_backbone.py | 8 +++--- .../models/densenet/densenet_backbone_test.py | 27 +++++++++++++++---- 2 files changed, 27 insertions(+), 8 deletions(-) diff --git a/keras_nlp/src/models/densenet/densenet_backbone.py b/keras_nlp/src/models/densenet/densenet_backbone.py index 60a5b28849..0ef700c5d6 100644 --- a/keras_nlp/src/models/densenet/densenet_backbone.py +++ b/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( @@ -103,7 +104,7 @@ def __init__( growth_rate, name=f"conv{len(stackwise_num_repeats) + 1}", ) - + pyramid_outputs[f"P{stack_index}"] = x x = keras.layers.BatchNormalization( axis=BN_AXIS, epsilon=BN_EPSILON, name="bn" )(x) @@ -117,6 +118,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() diff --git a/keras_nlp/src/models/densenet/densenet_backbone_test.py b/keras_nlp/src/models/densenet/densenet_backbone_test.py index 63f358035c..cd096982e3 100644 --- a/keras_nlp/src/models/densenet/densenet_backbone_test.py +++ b/keras_nlp/src/models/densenet/densenet_backbone_test.py @@ -14,6 +14,7 @@ import numpy as np import pytest +from keras import models from keras_nlp.src.models.densenet.densenet_backbone import DenseNetBackbone from keras_nlp.src.tests.test_case import TestCase @@ -22,23 +23,39 @@ 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): + init_kwargs = self.init_kwargs + backbone = DenseNetBackbone(**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)) + @pytest.mark.large def test_saved_model(self): self.run_model_saving_test( From 782d005a75e5f5df9ce35875242efa731f48b6b7 Mon Sep 17 00:00:00 2001 From: Sachin Prasad Date: Tue, 27 Aug 2024 22:34:00 +0000 Subject: [PATCH 2/7] fix testcase --- keras_nlp/src/models/densenet/densenet_backbone.py | 3 ++- keras_nlp/src/models/densenet/densenet_backbone_test.py | 2 +- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/keras_nlp/src/models/densenet/densenet_backbone.py b/keras_nlp/src/models/densenet/densenet_backbone.py index 0ef700c5d6..80a2122d0b 100644 --- a/keras_nlp/src/models/densenet/densenet_backbone.py +++ b/keras_nlp/src/models/densenet/densenet_backbone.py @@ -94,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}" ) @@ -104,7 +105,7 @@ def __init__( growth_rate, name=f"conv{len(stackwise_num_repeats) + 1}", ) - pyramid_outputs[f"P{stack_index}"] = x + pyramid_outputs[f"P{len(stackwise_num_repeats) +1}"] = x x = keras.layers.BatchNormalization( axis=BN_AXIS, epsilon=BN_EPSILON, name="bn" )(x) diff --git a/keras_nlp/src/models/densenet/densenet_backbone_test.py b/keras_nlp/src/models/densenet/densenet_backbone_test.py index cd096982e3..3108310ff9 100644 --- a/keras_nlp/src/models/densenet/densenet_backbone_test.py +++ b/keras_nlp/src/models/densenet/densenet_backbone_test.py @@ -51,7 +51,7 @@ def test_pyramid_output_format(self): self.assertEqual( list(output_data.keys()), list(backbone.pyramid_outputs.keys()) ) - self.assertEqual(list(output_data.keys()), ["P2", "P3", "P4"]) + self.assertEqual(list(output_data.keys()), ["P2", "P3", "P4", "P5"]) for k, v in output_data.items(): size = self.input_size // (2 ** int(k[1:])) self.assertEqual(tuple(v.shape[:3]), (2, size, size)) From 54e653240825424bddc4e8d6f1a8b3f5662b74cc Mon Sep 17 00:00:00 2001 From: Sachin Prasad Date: Tue, 27 Aug 2024 22:52:35 +0000 Subject: [PATCH 3/7] format fix --- keras_nlp/src/models/densenet/densenet_backbone_test.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/keras_nlp/src/models/densenet/densenet_backbone_test.py b/keras_nlp/src/models/densenet/densenet_backbone_test.py index 3108310ff9..2bcba0961f 100644 --- a/keras_nlp/src/models/densenet/densenet_backbone_test.py +++ b/keras_nlp/src/models/densenet/densenet_backbone_test.py @@ -30,7 +30,9 @@ def setUp(self): "image_shape": (32, 32, 3), } self.input_size = 32 - self.input_data = np.ones((2, self.input_size, self.input_size, 3), dtype="float32") + self.input_data = np.ones( + (2, self.input_size, self.input_size, 3), dtype="float32" + ) def test_backbone_basics(self): self.run_backbone_test( @@ -55,7 +57,7 @@ def test_pyramid_output_format(self): for k, v in output_data.items(): size = self.input_size // (2 ** int(k[1:])) self.assertEqual(tuple(v.shape[:3]), (2, size, size)) - + @pytest.mark.large def test_saved_model(self): self.run_model_saving_test( From bff4ffa6a7a2b9b55ba1d4cfc090dc8707cca11d Mon Sep 17 00:00:00 2001 From: Sachin Prasad Date: Wed, 28 Aug 2024 00:00:42 +0000 Subject: [PATCH 4/7] make common testcase for pyramid outputs --- .../csp_darknet/csp_darknet_backbone.py | 10 ++++++--- .../csp_darknet/csp_darknet_backbone_test.py | 19 ++++++++++++---- .../models/densenet/densenet_backbone_test.py | 18 +++++---------- .../src/models/resnet/resnet_backbone_test.py | 17 +++++--------- keras_nlp/src/tests/test_case.py | 22 +++++++++++++++++++ 5 files changed, 54 insertions(+), 32 deletions(-) diff --git a/keras_nlp/src/models/csp_darknet/csp_darknet_backbone.py b/keras_nlp/src/models/csp_darknet/csp_darknet_backbone.py index 607c6895ba..d5e4440e4c 100644 --- a/keras_nlp/src/models/csp_darknet/csp_darknet_backbone.py +++ b/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() diff --git a/keras_nlp/src/models/csp_darknet/csp_darknet_backbone_test.py b/keras_nlp/src/models/csp_darknet/csp_darknet_backbone_test.py index 857e06039d..8bbe3ba229 100644 --- a/keras_nlp/src/models/csp_darknet/csp_darknet_backbone_test.py +++ b/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( diff --git a/keras_nlp/src/models/densenet/densenet_backbone_test.py b/keras_nlp/src/models/densenet/densenet_backbone_test.py index 2bcba0961f..0b6b92879e 100644 --- a/keras_nlp/src/models/densenet/densenet_backbone_test.py +++ b/keras_nlp/src/models/densenet/densenet_backbone_test.py @@ -14,7 +14,6 @@ import numpy as np import pytest -from keras import models from keras_nlp.src.models.densenet.densenet_backbone import DenseNetBackbone from keras_nlp.src.tests.test_case import TestCase @@ -44,19 +43,12 @@ def test_backbone_basics(self): ) def test_pyramid_output_format(self): - init_kwargs = self.init_kwargs - backbone = DenseNetBackbone(**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=DenseNetBackbone, + init_kwargs=self.init_kwargs, + input_data=self.input_data, + expected_pyramid_output_keys=["P2", "P3", "P4", "P5"], ) - self.assertEqual(list(output_data.keys()), ["P2", "P3", "P4", "P5"]) - for k, v in output_data.items(): - size = self.input_size // (2 ** int(k[1:])) - self.assertEqual(tuple(v.shape[:3]), (2, size, size)) @pytest.mark.large def test_saved_model(self): diff --git a/keras_nlp/src/models/resnet/resnet_backbone_test.py b/keras_nlp/src/models/resnet/resnet_backbone_test.py index a6a30362cd..1a290fef73 100644 --- a/keras_nlp/src/models/resnet/resnet_backbone_test.py +++ b/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"), diff --git a/keras_nlp/src/tests/test_case.py b/keras_nlp/src/tests/test_case.py index 8e63bc19d9..9a3dc9636a 100644 --- a/keras_nlp/src/tests/test_case.py +++ b/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") From feca564826445e59aab3148c5efc8cda0feb7dbe Mon Sep 17 00:00:00 2001 From: sachin prasad Date: Thu, 29 Aug 2024 11:28:16 -0700 Subject: [PATCH 5/7] change default shape --- keras_nlp/src/models/csp_darknet/csp_darknet_backbone.py | 4 ++-- keras_nlp/src/models/densenet/densenet_backbone.py | 4 ++-- .../src/models/mix_transformer/mix_transformer_backbone.py | 4 ++-- 3 files changed, 6 insertions(+), 6 deletions(-) diff --git a/keras_nlp/src/models/csp_darknet/csp_darknet_backbone.py b/keras_nlp/src/models/csp_darknet/csp_darknet_backbone.py index d5e4440e4c..9f7b515cc9 100644 --- a/keras_nlp/src/models/csp_darknet/csp_darknet_backbone.py +++ b/keras_nlp/src/models/csp_darknet/csp_darknet_backbone.py @@ -39,7 +39,7 @@ class CSPDarkNetBackbone(FeaturePyramidBackbone): `"basic_block"` for basic conv block. Defaults to "basic_block". image_shape: tuple. The input shape without the batch size. - Defaults to `(224, 224, 3)`. + Defaults to `(None, None, 3)`. Examples: ```python @@ -67,7 +67,7 @@ def __init__( stackwise_depth, include_rescaling, block_type="basic_block", - image_shape=(224, 224, 3), + image_shape=(None, None, 3), **kwargs, ): # === Functional Model === diff --git a/keras_nlp/src/models/densenet/densenet_backbone.py b/keras_nlp/src/models/densenet/densenet_backbone.py index 80a2122d0b..48940718a1 100644 --- a/keras_nlp/src/models/densenet/densenet_backbone.py +++ b/keras_nlp/src/models/densenet/densenet_backbone.py @@ -35,7 +35,7 @@ class DenseNetBackbone(FeaturePyramidBackbone): include_rescaling: bool, whether to rescale the inputs. If set to `True`, inputs will be passed through a `Rescaling(1/255.0)` layer. Defaults to `True`. - image_shape: optional shape tuple, defaults to (224, 224, 3). + image_shape: optional shape tuple, defaults to (None, None, 3). compression_ratio: float, compression rate at transition layers, defaults to 0.5. growth_rate: int, number of filters added by each dense block, @@ -62,7 +62,7 @@ def __init__( self, stackwise_num_repeats, include_rescaling=True, - image_shape=(224, 224, 3), + image_shape=(None, None, 3), compression_ratio=0.5, growth_rate=32, **kwargs, diff --git a/keras_nlp/src/models/mix_transformer/mix_transformer_backbone.py b/keras_nlp/src/models/mix_transformer/mix_transformer_backbone.py index 2cfe7f6761..35c5f7fd5a 100644 --- a/keras_nlp/src/models/mix_transformer/mix_transformer_backbone.py +++ b/keras_nlp/src/models/mix_transformer/mix_transformer_backbone.py @@ -37,7 +37,7 @@ def __init__( patch_sizes, strides, include_rescaling=True, - image_shape=(224, 224, 3), + image_shape=(None, None, 3), hidden_dims=None, **kwargs, ): @@ -63,7 +63,7 @@ def __init__( include_rescaling: bool, whether to rescale the inputs. If set to `True`, inputs will be passed through a `Rescaling(1/255.0)` layer. Defaults to `True`. - image_shape: optional shape tuple, defaults to (224, 224, 3). + image_shape: optional shape tuple, defaults to (None, None, 3). hidden_dims: the embedding dims per hierarchical layer, used as the levels of the feature pyramid. patch_sizes: list of integers, the patch_size to apply for each layer. From a4ef81a2744b40ac3310a37bba3bfe3fcf15e5a8 Mon Sep 17 00:00:00 2001 From: sachin prasad Date: Thu, 29 Aug 2024 14:29:04 -0700 Subject: [PATCH 6/7] simplify testcase --- .../csp_darknet/csp_darknet_backbone_test.py | 11 ++--------- .../src/models/densenet/densenet_backbone_test.py | 11 ++--------- .../src/models/resnet/resnet_backbone_test.py | 11 ----------- keras_nlp/src/tests/test_case.py | 15 +++++++++++++++ 4 files changed, 19 insertions(+), 29 deletions(-) diff --git a/keras_nlp/src/models/csp_darknet/csp_darknet_backbone_test.py b/keras_nlp/src/models/csp_darknet/csp_darknet_backbone_test.py index 8bbe3ba229..c8efd406a0 100644 --- a/keras_nlp/src/models/csp_darknet/csp_darknet_backbone_test.py +++ b/keras_nlp/src/models/csp_darknet/csp_darknet_backbone_test.py @@ -36,20 +36,13 @@ def setUp(self): ) def test_backbone_basics(self): - self.run_backbone_test( + self.run_vision_backbone_test( cls=CSPDarkNetBackbone, init_kwargs=self.init_kwargs, input_data=self.input_data, 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"], + run_mixed_precision_check=False, ) @pytest.mark.large diff --git a/keras_nlp/src/models/densenet/densenet_backbone_test.py b/keras_nlp/src/models/densenet/densenet_backbone_test.py index 0b6b92879e..35c5850c18 100644 --- a/keras_nlp/src/models/densenet/densenet_backbone_test.py +++ b/keras_nlp/src/models/densenet/densenet_backbone_test.py @@ -34,20 +34,13 @@ def setUp(self): ) def test_backbone_basics(self): - self.run_backbone_test( + self.run_vision_backbone_test( cls=DenseNetBackbone, init_kwargs=self.init_kwargs, input_data=self.input_data, 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"], + run_mixed_precision_check=False, ) @pytest.mark.large diff --git a/keras_nlp/src/models/resnet/resnet_backbone_test.py b/keras_nlp/src/models/resnet/resnet_backbone_test.py index 1a290fef73..f97b854d51 100644 --- a/keras_nlp/src/models/resnet/resnet_backbone_test.py +++ b/keras_nlp/src/models/resnet/resnet_backbone_test.py @@ -50,17 +50,6 @@ def test_backbone_basics(self, use_pre_activation, block_type): expected_output_shape=( (2, 64) if block_type == "basic_block" else (2, 256) ), - ) - - def test_pyramid_output_format(self): - init_kwargs = self.init_kwargs.copy() - init_kwargs.update( - {"block_type": "basic_block", "use_pre_activation": False} - ) - self.run_pyramid_output_test( - cls=ResNetBackbone, - init_kwargs=init_kwargs, - input_data=self.input_data, expected_pyramid_output_keys=["P2", "P3", "P4"], ) diff --git a/keras_nlp/src/tests/test_case.py b/keras_nlp/src/tests/test_case.py index 9a3dc9636a..a2c89be90e 100644 --- a/keras_nlp/src/tests/test_case.py +++ b/keras_nlp/src/tests/test_case.py @@ -465,6 +465,7 @@ def run_vision_backbone_test( init_kwargs, input_data, expected_output_shape, + expected_pyramid_output_keys=None, variable_length_data=None, run_mixed_precision_check=True, run_quantization_check=True, @@ -491,6 +492,13 @@ def run_vision_backbone_test( run_mixed_precision_check=run_mixed_precision_check, run_quantization_check=run_quantization_check, ) + if expected_pyramid_output_keys: + self.run_pyramid_output_test( + cls=cls, + init_kwargs=init_kwargs, + input_data=input_data, + expected_pyramid_output_keys=expected_pyramid_output_keys, + ) # Check data_format. We assume that `input_data` is in "channels_last" # format. @@ -515,6 +523,13 @@ def run_vision_backbone_test( run_mixed_precision_check=run_mixed_precision_check, run_quantization_check=run_quantization_check, ) + if expected_pyramid_output_keys: + self.run_pyramid_output_test( + cls=cls, + init_kwargs=init_kwargs, + input_data=input_data, + expected_pyramid_output_keys=expected_pyramid_output_keys, + ) # Restore the original `image_data_format`. keras.config.set_image_data_format(ori_data_format) From 8e65eb12be3abbde48c78768c7e466106439e676 Mon Sep 17 00:00:00 2001 From: sachin prasad Date: Fri, 30 Aug 2024 15:52:54 -0700 Subject: [PATCH 7/7] test case change and add channel axis --- .../csp_darknet/csp_darknet_backbone.py | 59 +++++++++++++++---- .../csp_darknet/csp_darknet_backbone_test.py | 3 +- .../src/models/densenet/densenet_backbone.py | 55 +++++++++++------ .../models/densenet/densenet_backbone_test.py | 3 +- .../mix_transformer_backbone_test.py | 21 ++----- .../src/models/resnet/resnet_backbone_test.py | 1 + keras_nlp/src/tests/test_case.py | 52 ++++++---------- 7 files changed, 112 insertions(+), 82 deletions(-) diff --git a/keras_nlp/src/models/csp_darknet/csp_darknet_backbone.py b/keras_nlp/src/models/csp_darknet/csp_darknet_backbone.py index 9f7b515cc9..40efb6de04 100644 --- a/keras_nlp/src/models/csp_darknet/csp_darknet_backbone.py +++ b/keras_nlp/src/models/csp_darknet/csp_darknet_backbone.py @@ -65,12 +65,15 @@ def __init__( self, stackwise_num_filters, stackwise_depth, - include_rescaling, + include_rescaling=True, block_type="basic_block", image_shape=(None, None, 3), **kwargs, ): # === Functional Model === + channel_axis = ( + -1 if keras.config.image_data_format() == "channels_last" else 1 + ) apply_ConvBlock = ( apply_darknet_conv_block_depthwise if block_type == "depthwise_block" @@ -83,9 +86,13 @@ def __init__( if include_rescaling: x = layers.Rescaling(scale=1 / 255.0)(x) - x = apply_focus(name="stem_focus")(x) + x = apply_focus(channel_axis, name="stem_focus")(x) x = apply_darknet_conv_block( - base_channels, kernel_size=3, strides=1, name="stem_conv" + base_channels, + channel_axis, + kernel_size=3, + strides=1, + name="stem_conv", )(x) pyramid_outputs = {} @@ -94,6 +101,7 @@ def __init__( ): x = apply_ConvBlock( channels, + channel_axis, kernel_size=3, strides=2, name=f"dark{index + 2}_conv", @@ -102,12 +110,14 @@ def __init__( if index == len(stackwise_depth) - 1: x = apply_spatial_pyramid_pooling_bottleneck( channels, + channel_axis, hidden_filters=channels // 2, name=f"dark{index + 2}_spp", )(x) x = apply_cross_stage_partial( channels, + channel_axis, num_bottlenecks=depth, block_type="basic_block", residual=(index != len(stackwise_depth) - 1), @@ -139,7 +149,7 @@ def get_config(self): return config -def apply_focus(name=None): +def apply_focus(channel_axis, name=None): """A block used in CSPDarknet to focus information into channels of the image. @@ -155,7 +165,7 @@ def apply_focus(name=None): """ def apply(x): - return layers.Concatenate(name=name)( + return layers.Concatenate(axis=channel_axis, name=name)( [ x[..., ::2, ::2, :], x[..., 1::2, ::2, :], @@ -168,7 +178,13 @@ def apply(x): def apply_darknet_conv_block( - filters, kernel_size, strides, use_bias=False, activation="silu", name=None + filters, + channel_axis, + kernel_size, + strides, + use_bias=False, + activation="silu", + name=None, ): """ The basic conv block used in Darknet. Applies Conv2D followed by a @@ -197,11 +213,12 @@ def apply(inputs): kernel_size, strides, padding="same", + data_format=keras.config.image_data_format(), use_bias=use_bias, name=name + "_conv", )(inputs) - x = layers.BatchNormalization(name=name + "_bn")(x) + x = layers.BatchNormalization(axis=channel_axis, name=name + "_bn")(x) if activation == "silu": x = layers.Lambda(lambda x: keras.activations.silu(x))(x) @@ -216,7 +233,7 @@ def apply(inputs): def apply_darknet_conv_block_depthwise( - filters, kernel_size, strides, activation="silu", name=None + filters, channel_axis, kernel_size, strides, activation="silu", name=None ): """ The depthwise conv block used in CSPDarknet. @@ -240,9 +257,13 @@ def apply_darknet_conv_block_depthwise( def apply(inputs): x = layers.DepthwiseConv2D( - kernel_size, strides, padding="same", use_bias=False + kernel_size, + strides, + padding="same", + data_format=keras.config.image_data_format(), + use_bias=False, )(inputs) - x = layers.BatchNormalization()(x) + x = layers.BatchNormalization(axis=channel_axis)(x) if activation == "silu": x = layers.Lambda(lambda x: keras.activations.swish(x))(x) @@ -252,7 +273,11 @@ def apply(inputs): x = layers.LeakyReLU(0.1)(x) x = apply_darknet_conv_block( - filters, kernel_size=1, strides=1, activation=activation + filters, + channel_axis, + kernel_size=1, + strides=1, + activation=activation, )(x) return x @@ -262,6 +287,7 @@ def apply(inputs): def apply_spatial_pyramid_pooling_bottleneck( filters, + channel_axis, hidden_filters=None, kernel_sizes=(5, 9, 13), activation="silu", @@ -295,6 +321,7 @@ def apply_spatial_pyramid_pooling_bottleneck( def apply(x): x = apply_darknet_conv_block( hidden_filters, + channel_axis, kernel_size=1, strides=1, activation=activation, @@ -308,13 +335,15 @@ def apply(x): kernel_size, strides=1, padding="same", + data_format=keras.config.image_data_format(), name=f"{name}_maxpool_{kernel_size}", )(x[0]) ) - x = layers.Concatenate(name=f"{name}_concat")(x) + x = layers.Concatenate(axis=channel_axis, name=f"{name}_concat")(x) x = apply_darknet_conv_block( filters, + channel_axis, kernel_size=1, strides=1, activation=activation, @@ -328,6 +357,7 @@ def apply(x): def apply_cross_stage_partial( filters, + channel_axis, num_bottlenecks, residual=True, block_type="basic_block", @@ -365,6 +395,7 @@ def apply(inputs): x1 = apply_darknet_conv_block( hidden_channels, + channel_axis, kernel_size=1, strides=1, activation=activation, @@ -373,6 +404,7 @@ def apply(inputs): x2 = apply_darknet_conv_block( hidden_channels, + channel_axis, kernel_size=1, strides=1, activation=activation, @@ -383,6 +415,7 @@ def apply(inputs): residual_x = x1 x1 = apply_darknet_conv_block( hidden_channels, + channel_axis, kernel_size=1, strides=1, activation=activation, @@ -390,6 +423,7 @@ def apply(inputs): )(x1) x1 = ConvBlock( hidden_channels, + channel_axis, kernel_size=3, strides=1, activation=activation, @@ -403,6 +437,7 @@ def apply(inputs): x = layers.Concatenate(name=f"{name}_concat")([x1, x2]) x = apply_darknet_conv_block( filters, + channel_axis, kernel_size=1, strides=1, activation=activation, diff --git a/keras_nlp/src/models/csp_darknet/csp_darknet_backbone_test.py b/keras_nlp/src/models/csp_darknet/csp_darknet_backbone_test.py index c8efd406a0..ed6dc7b525 100644 --- a/keras_nlp/src/models/csp_darknet/csp_darknet_backbone_test.py +++ b/keras_nlp/src/models/csp_darknet/csp_darknet_backbone_test.py @@ -26,7 +26,6 @@ def setUp(self): self.init_kwargs = { "stackwise_num_filters": [2, 4, 6, 8], "stackwise_depth": [1, 3, 3, 1], - "include_rescaling": False, "block_type": "basic_block", "image_shape": (32, 32, 3), } @@ -42,7 +41,9 @@ def test_backbone_basics(self): input_data=self.input_data, expected_output_shape=(2, 1, 1, 8), expected_pyramid_output_keys=["P2", "P3", "P4", "P5"], + expected_pyramid_image_sizes=[(8, 8), (4, 4), (2, 2), (1, 1)], run_mixed_precision_check=False, + run_data_format_check=False, ) @pytest.mark.large diff --git a/keras_nlp/src/models/densenet/densenet_backbone.py b/keras_nlp/src/models/densenet/densenet_backbone.py index 48940718a1..13e3d8597f 100644 --- a/keras_nlp/src/models/densenet/densenet_backbone.py +++ b/keras_nlp/src/models/densenet/densenet_backbone.py @@ -16,7 +16,6 @@ from keras_nlp.src.api_export import keras_nlp_export from keras_nlp.src.models.feature_pyramid_backbone import FeaturePyramidBackbone -BN_AXIS = 3 BN_EPSILON = 1.001e-5 @@ -68,6 +67,8 @@ def __init__( **kwargs, ): # === Functional Model === + data_format = keras.config.image_data_format() + channel_axis = -1 if data_format == "channels_last" else 1 image_input = keras.layers.Input(shape=image_shape) x = image_input @@ -75,14 +76,20 @@ def __init__( x = keras.layers.Rescaling(1 / 255.0)(x) x = keras.layers.Conv2D( - 64, 7, strides=2, use_bias=False, padding="same", name="conv1_conv" + 64, + 7, + strides=2, + use_bias=False, + padding="same", + data_format=data_format, + name="conv1_conv", )(x) x = keras.layers.BatchNormalization( - axis=BN_AXIS, epsilon=BN_EPSILON, name="conv1_bn" + axis=channel_axis, epsilon=BN_EPSILON, name="conv1_bn" )(x) x = keras.layers.Activation("relu", name="conv1_relu")(x) x = keras.layers.MaxPooling2D( - 3, strides=2, padding="same", name="pool1" + 3, strides=2, padding="same", data_format=data_format, name="pool1" )(x) pyramid_outputs = {} @@ -90,24 +97,26 @@ def __init__( index = stack_index + 2 x = apply_dense_block( x, + channel_axis, stackwise_num_repeats[stack_index], growth_rate, name=f"conv{index}", ) pyramid_outputs[f"P{index}"] = x x = apply_transition_block( - x, compression_ratio, name=f"pool{index}" + x, channel_axis, compression_ratio, name=f"pool{index}" ) x = apply_dense_block( x, + channel_axis, stackwise_num_repeats[-1], 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" + axis=channel_axis, epsilon=BN_EPSILON, name="bn" )(x) x = keras.layers.Activation("relu", name="relu")(x) @@ -135,7 +144,7 @@ def get_config(self): return config -def apply_dense_block(x, num_repeats, growth_rate, name=None): +def apply_dense_block(x, channel_axis, num_repeats, growth_rate, name=None): """A dense block. Args: @@ -148,11 +157,13 @@ def apply_dense_block(x, num_repeats, growth_rate, name=None): name = f"dense_block_{keras.backend.get_uid('dense_block')}" for i in range(num_repeats): - x = apply_conv_block(x, growth_rate, name=f"{name}_block_{i}") + x = apply_conv_block( + x, channel_axis, growth_rate, name=f"{name}_block_{i}" + ) return x -def apply_transition_block(x, compression_ratio, name=None): +def apply_transition_block(x, channel_axis, compression_ratio, name=None): """A transition block. Args: @@ -160,24 +171,28 @@ def apply_transition_block(x, compression_ratio, name=None): compression_ratio: float, compression rate at transition layers. name: string, block label. """ + data_format = keras.config.image_data_format() if name is None: name = f"transition_block_{keras.backend.get_uid('transition_block')}" x = keras.layers.BatchNormalization( - axis=BN_AXIS, epsilon=BN_EPSILON, name=f"{name}_bn" + axis=channel_axis, epsilon=BN_EPSILON, name=f"{name}_bn" )(x) x = keras.layers.Activation("relu", name=f"{name}_relu")(x) x = keras.layers.Conv2D( - int(x.shape[BN_AXIS] * compression_ratio), + int(x.shape[channel_axis] * compression_ratio), 1, use_bias=False, + data_format=data_format, name=f"{name}_conv", )(x) - x = keras.layers.AveragePooling2D(2, strides=2, name=f"{name}_pool")(x) + x = keras.layers.AveragePooling2D( + 2, strides=2, data_format=data_format, name=f"{name}_pool" + )(x) return x -def apply_conv_block(x, growth_rate, name=None): +def apply_conv_block(x, channel_axis, growth_rate, name=None): """A building block for a dense block. Args: @@ -185,19 +200,24 @@ def apply_conv_block(x, growth_rate, name=None): growth_rate: int, number of filters added by each dense block. name: string, block label. """ + data_format = keras.config.image_data_format() if name is None: name = f"conv_block_{keras.backend.get_uid('conv_block')}" shortcut = x x = keras.layers.BatchNormalization( - axis=BN_AXIS, epsilon=BN_EPSILON, name=f"{name}_0_bn" + axis=channel_axis, epsilon=BN_EPSILON, name=f"{name}_0_bn" )(x) x = keras.layers.Activation("relu", name=f"{name}_0_relu")(x) x = keras.layers.Conv2D( - 4 * growth_rate, 1, use_bias=False, name=f"{name}_1_conv" + 4 * growth_rate, + 1, + use_bias=False, + data_format=data_format, + name=f"{name}_1_conv", )(x) x = keras.layers.BatchNormalization( - axis=BN_AXIS, epsilon=BN_EPSILON, name=f"{name}_1_bn" + axis=channel_axis, epsilon=BN_EPSILON, name=f"{name}_1_bn" )(x) x = keras.layers.Activation("relu", name=f"{name}_1_relu")(x) x = keras.layers.Conv2D( @@ -205,9 +225,10 @@ def apply_conv_block(x, growth_rate, name=None): 3, padding="same", use_bias=False, + data_format=data_format, name=f"{name}_2_conv", )(x) - x = keras.layers.Concatenate(axis=BN_AXIS, name=f"{name}_concat")( + x = keras.layers.Concatenate(axis=channel_axis, name=f"{name}_concat")( [shortcut, x] ) return x diff --git a/keras_nlp/src/models/densenet/densenet_backbone_test.py b/keras_nlp/src/models/densenet/densenet_backbone_test.py index 35c5850c18..7720411319 100644 --- a/keras_nlp/src/models/densenet/densenet_backbone_test.py +++ b/keras_nlp/src/models/densenet/densenet_backbone_test.py @@ -23,7 +23,6 @@ class DenseNetBackboneTest(TestCase): def setUp(self): self.init_kwargs = { "stackwise_num_repeats": [2, 4, 6, 4], - "include_rescaling": True, "compression_ratio": 0.5, "growth_rate": 2, "image_shape": (32, 32, 3), @@ -40,7 +39,9 @@ def test_backbone_basics(self): input_data=self.input_data, expected_output_shape=(2, 1, 1, 24), expected_pyramid_output_keys=["P2", "P3", "P4", "P5"], + expected_pyramid_image_sizes=[(8, 8), (4, 4), (2, 2), (1, 1)], run_mixed_precision_check=False, + run_data_format_check=False, ) @pytest.mark.large diff --git a/keras_nlp/src/models/mix_transformer/mix_transformer_backbone_test.py b/keras_nlp/src/models/mix_transformer/mix_transformer_backbone_test.py index 4f1955297f..280adca065 100644 --- a/keras_nlp/src/models/mix_transformer/mix_transformer_backbone_test.py +++ b/keras_nlp/src/models/mix_transformer/mix_transformer_backbone_test.py @@ -14,7 +14,6 @@ import numpy as np import pytest -from keras import models from keras_nlp.src.models.mix_transformer.mix_transformer_backbone import ( MiTBackbone, @@ -42,30 +41,18 @@ def setUp(self): ) def test_backbone_basics(self): - self.run_backbone_test( + self.run_vision_backbone_test( cls=MiTBackbone, init_kwargs=self.init_kwargs, input_data=self.input_data, expected_output_shape=(2, 2, 2, 8), + expected_pyramid_output_keys=["P1", "P2"], + expected_pyramid_image_sizes=[(4, 4), (2, 2)], run_quantization_check=False, run_mixed_precision_check=False, + run_data_format_check=False, ) - def test_pyramid_output_format(self): - init_kwargs = self.init_kwargs - backbone = MiTBackbone(**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()), ["P1", "P2"]) - for k, v in output_data.items(): - size = self.input_size // (2 ** (int(k[1:]) + 1)) - self.assertEqual(tuple(v.shape[:3]), (2, size, size)) - @pytest.mark.large def test_saved_model(self): self.run_model_saving_test( diff --git a/keras_nlp/src/models/resnet/resnet_backbone_test.py b/keras_nlp/src/models/resnet/resnet_backbone_test.py index f97b854d51..fdf52cab73 100644 --- a/keras_nlp/src/models/resnet/resnet_backbone_test.py +++ b/keras_nlp/src/models/resnet/resnet_backbone_test.py @@ -51,6 +51,7 @@ def test_backbone_basics(self, use_pre_activation, block_type): (2, 64) if block_type == "basic_block" else (2, 256) ), expected_pyramid_output_keys=["P2", "P3", "P4"], + expected_pyramid_image_sizes=[(16, 16), (8, 8), (4, 4)], ) @parameterized.named_parameters( diff --git a/keras_nlp/src/tests/test_case.py b/keras_nlp/src/tests/test_case.py index a2c89be90e..6dad75b84d 100644 --- a/keras_nlp/src/tests/test_case.py +++ b/keras_nlp/src/tests/test_case.py @@ -466,6 +466,7 @@ def run_vision_backbone_test( input_data, expected_output_shape, expected_pyramid_output_keys=None, + expected_pyramid_image_sizes=None, variable_length_data=None, run_mixed_precision_check=True, run_quantization_check=True, @@ -493,12 +494,24 @@ def run_vision_backbone_test( run_quantization_check=run_quantization_check, ) if expected_pyramid_output_keys: - self.run_pyramid_output_test( - cls=cls, - init_kwargs=init_kwargs, - input_data=input_data, - expected_pyramid_output_keys=expected_pyramid_output_keys, + 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 ) + # check height and width of each level. + for i, (k, v) in enumerate(output_data.items()): + self.assertEqual( + tuple(v.shape[1:3]), expected_pyramid_image_sizes[i] + ) # Check data_format. We assume that `input_data` is in "channels_last" # format. @@ -523,13 +536,6 @@ def run_vision_backbone_test( run_mixed_precision_check=run_mixed_precision_check, run_quantization_check=run_quantization_check, ) - if expected_pyramid_output_keys: - self.run_pyramid_output_test( - cls=cls, - init_kwargs=init_kwargs, - input_data=input_data, - expected_pyramid_output_keys=expected_pyramid_output_keys, - ) # Restore the original `image_data_format`. keras.config.set_image_data_format(ori_data_format) @@ -619,27 +625,5 @@ 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")