Merge pull request keras-team#4 from ianstenbit/core-csp

Port CSPDarkNet and DarkNet to Keras Core

Author: ianstenbit

.github/workflows/actions.yml

@@ -85,6 +85,7 @@ jobs:
         pytest keras_cv/models/backbones/resnet_v1 \
                keras_cv/models/backbones/resnet_v2 \
                keras_cv/models/backbones/efficientnet_v2 \
+               keras_cv/models/backbones/csp_darknet \
                keras_cv/models/classification/image_classifier_test.py \
                --durations 0 --run_large --run_extra_large
   format:

keras_cv/backend/__init__.py

@@ -61,4 +61,7 @@
 else:
     from tensorflow import keras
 
+    # TF Keras doesn't have this rename.
+    keras.activations.silu = keras.activations.swish
+
 from keras_cv.backend import ops  # noqa: E402

keras_cv/models/backbones/csp_darknet/csp_darknet_backbone.py

@@ -15,9 +15,7 @@
 """CSPDarkNet models for KerasCV. """
 import copy
 
-from tensorflow import keras
-from tensorflow.keras import layers
-
+from keras_cv.backend import keras
 from keras_cv.models import utils
 from keras_cv.models.backbones.backbone import Backbone
 from keras_cv.models.backbones.csp_darknet.csp_darknet_backbone_presets import (
@@ -64,8 +62,8 @@ class CSPDarkNetBackbone(Backbone):
         use_depthwise: bool, whether a `DarknetConvBlockDepthwise` should be
             used over a `DarknetConvBlock`, defaults to False.
         input_shape: optional shape tuple, defaults to (None, None, 3).
-        input_tensor: optional Keras tensor (i.e. output of `layers.Input()`)
-            to use as image input for the model.
+        input_tensor: optional Keras tensor (i.e. output of
+            `keras.layers.Input()`) to use as image input for the model.
 
     Returns:
         A `keras.Model` instance.
@@ -111,7 +109,7 @@ def __init__(
 
         x = inputs
         if include_rescaling:
-            x = layers.Rescaling(1 / 255.0)(x)
+            x = keras.layers.Rescaling(1 / 255.0)(x)
 
         # stem
         x = Focus(name="stem_focus")(x)
@@ -144,7 +142,9 @@ def __init__(
                 residual=(index != len(stackwise_depth) - 1),
                 name=f"dark{index + 2}_csp",
             )(x)
-            pyramid_level_inputs[f"P{index + 2}"] = x.node.layer.name
+            pyramid_level_inputs[f"P{index + 2}"] = utils.get_tensor_input_name(
+                x
+            )
 
         super().__init__(inputs=inputs, outputs=x, **kwargs)
         self.pyramid_level_inputs = pyramid_level_inputs
@@ -195,8 +195,8 @@ def presets_with_weights(cls):
     Args:
         include_rescaling: bool, whether or not to rescale the inputs. If set to
             True, inputs will be passed through a `Rescaling(1/255.0)` layer.
-        input_tensor: optional Keras tensor (i.e. output of `layers.Input()`)
-            to use as image input for the model.
+        input_tensor: optional Keras tensor (i.e. output of
+            `keras.layers.Input()`) to use as image input for the model.
         input_shape: optional shape tuple, defaults to (None, None, 3).
 
     Examples:

keras_cv/models/backbones/csp_darknet/csp_darknet_backbone_test.py

@@ -14,18 +14,19 @@
 
 import os
 
+import numpy as np
 import pytest
 import tensorflow as tf
 from absl.testing import parameterized
-from tensorflow import keras
 
+from keras_cv.backend import keras
 from keras_cv.models.backbones.csp_darknet import csp_darknet_backbone
 from keras_cv.utils.train import get_feature_extractor
 
 
 class CSPDarkNetBackboneTest(tf.test.TestCase, parameterized.TestCase):
     def setUp(self):
-        self.input_batch = tf.ones(shape=(2, 224, 224, 3))
+        self.input_batch = np.ones(shape=(2, 224, 224, 3))
 
     def test_valid_call(self):
         model = csp_darknet_backbone.CSPDarkNetBackbone(
@@ -47,20 +48,18 @@ def test_valid_call_with_rescaling(self):
         )
         model(self.input_batch)
 
-    @parameterized.named_parameters(
-        ("tf_format", "tf", "model"),
-        ("keras_format", "keras_v3", "model.keras"),
-    )
     @pytest.mark.large  # Saving is slow, so mark these large.
-    def test_saved_model(self, save_format, filename):
+    def test_saved_model(self):
         model = csp_darknet_backbone.CSPDarkNetBackbone(
             stackwise_channels=[48, 96, 192, 384],
             stackwise_depth=[1, 3, 3, 1],
             include_rescaling=True,
         )
         model_output = model(self.input_batch)
-        save_path = os.path.join(self.get_temp_dir(), filename)
-        model.save(save_path, save_format=save_format)
+        save_path = os.path.join(
+            self.get_temp_dir(), "csp_darknet_backbone.keras"
+        )
+        model.save(save_path)
         restored_model = keras.models.load_model(save_path)
 
         # Check we got the real object back.
@@ -72,16 +71,14 @@ def test_saved_model(self, save_format, filename):
         restored_output = restored_model(self.input_batch)
         self.assertAllClose(model_output, restored_output)
 
-    @parameterized.named_parameters(
-        ("tf_format", "tf", "model"),
-        ("keras_format", "keras_v3", "model.keras"),
-    )
     @pytest.mark.large  # Saving is slow, so mark these large.
-    def test_saved_alias_model(self, save_format, filename):
+    def test_saved_alias_model(self):
         model = csp_darknet_backbone.CSPDarkNetLBackbone()
         model_output = model(self.input_batch)
-        save_path = os.path.join(self.get_temp_dir(), filename)
-        model.save(save_path, save_format=save_format)
+        save_path = os.path.join(
+            self.get_temp_dir(), "csp_darknet_backbone.keras"
+        )
+        model.save(save_path)
         restored_model = keras.models.load_model(save_path)
 
         # Check we got the real object back.
@@ -109,10 +106,10 @@ def test_create_backbone_model_from_alias_model(self):
         levels = ["P2", "P3", "P4", "P5"]
         self.assertLen(outputs, 4)
         self.assertEquals(list(outputs.keys()), levels)
-        self.assertEquals(outputs["P2"].shape, [None, 56, 56, 128])
-        self.assertEquals(outputs["P3"].shape, [None, 28, 28, 256])
-        self.assertEquals(outputs["P4"].shape, [None, 14, 14, 512])
-        self.assertEquals(outputs["P5"].shape, [None, 7, 7, 1024])
+        self.assertEquals(outputs["P2"].shape, (None, 56, 56, 128))
+        self.assertEquals(outputs["P3"].shape, (None, 28, 28, 256))
+        self.assertEquals(outputs["P4"].shape, (None, 14, 14, 512))
+        self.assertEquals(outputs["P5"].shape, (None, 7, 7, 1024))
 
     def test_create_backbone_model_with_level_config(self):
         model = csp_darknet_backbone.CSPDarkNetBackbone(
@@ -127,8 +124,8 @@ def test_create_backbone_model_with_level_config(self):
         outputs = backbone_model(inputs)
         self.assertLen(outputs, 2)
         self.assertEquals(list(outputs.keys()), levels)
-        self.assertEquals(outputs["P3"].shape, [None, 32, 32, 96])
-        self.assertEquals(outputs["P4"].shape, [None, 16, 16, 192])
+        self.assertEquals(outputs["P3"].shape, (None, 32, 32, 96))
+        self.assertEquals(outputs["P4"].shape, (None, 16, 16, 192))
 
     @parameterized.named_parameters(
         ("Tiny", csp_darknet_backbone.CSPDarkNetTinyBackbone),
@@ -139,7 +136,7 @@ def test_create_backbone_model_with_level_config(self):
     )
     def test_specific_arch_forward_pass(self, arch_class):
         backbone = arch_class()
-        backbone(tf.random.uniform(shape=[2, 256, 256, 3]))
+        backbone(np.random.uniform(size=(2, 256, 256, 3)))
 
     @parameterized.named_parameters(
         ("Tiny", csp_darknet_backbone.CSPDarkNetTinyBackbone),

keras_cv/models/backbones/csp_darknet/csp_darknet_utils.py

@@ -18,10 +18,7 @@
   - [YoloV3 implementation](https://github.com/ultralytics/yolov3)
 """
 
-import tensorflow as tf
-from tensorflow import keras
-from tensorflow.keras import backend
-from tensorflow.keras import layers
+from keras_cv.backend import keras
 
 
 def DarknetConvBlock(
@@ -46,28 +43,30 @@ def DarknetConvBlock(
     """
 
     if name is None:
-        name = f"conv_block{backend.get_uid('conv_block')}"
+        name = f"conv_block{keras.backend.get_uid('conv_block')}"
 
     model_layers = [
-        layers.Conv2D(
+        keras.layers.Conv2D(
             filters,
             kernel_size,
             strides,
             padding="same",
             use_bias=use_bias,
             name=name + "_conv",
         ),
-        layers.BatchNormalization(name=name + "_bn"),
+        keras.layers.BatchNormalization(name=name + "_bn"),
     ]
 
     if activation == "silu":
-        model_layers.append(layers.Lambda(lambda x: keras.activations.swish(x)))
+        model_layers.append(
+            keras.layers.Lambda(lambda x: keras.activations.silu(x))
+        )
     elif activation == "relu":
-        model_layers.append(layers.ReLU())
+        model_layers.append(keras.layers.ReLU())
     elif activation == "leaky_relu":
-        model_layers.append(layers.LeakyReLU(0.1))
+        model_layers.append(keras.layers.LeakyReLU(0.1))
 
-    return keras.Sequential(model_layers, name=None)
+    return keras.Sequential(model_layers, name=name)
 
 
 def ResidualBlocks(filters, num_blocks, name=None):
@@ -84,7 +83,7 @@ def ResidualBlocks(filters, num_blocks, name=None):
     """
 
     if name is None:
-        name = f"residual_block{backend.get_uid('residual_block')}"
+        name = f"residual_block{keras.backend.get_uid('residual_block')}"
 
     def apply(x):
         x = DarknetConvBlock(
@@ -114,9 +113,9 @@ def apply(x):
             )(x)
 
             if i == num_blocks:
-                x = layers.Add(name=f"{name}_out")([residual, x])
+                x = keras.layers.Add(name=f"{name}_out")([residual, x])
             else:
-                x = layers.Add(name=f"{name}_add_{i}")([residual, x])
+                x = keras.layers.Add(name=f"{name}_add_{i}")([residual, x])
 
         return x
 
@@ -149,7 +148,7 @@ def SpatialPyramidPoolingBottleneck(
         SpatialPyramidPoolingBottleneck.
     """
     if name is None:
-        name = f"spp{backend.get_uid('spp')}"
+        name = f"spp{keras.backend.get_uid('spp')}"
 
     if hidden_filters is None:
         hidden_filters = filters
@@ -166,15 +165,15 @@ def apply(x):
 
         for kernel_size in kernel_sizes:
             x.append(
-                layers.MaxPooling2D(
+                keras.layers.MaxPooling2D(
                     kernel_size,
                     strides=1,
                     padding="same",
                     name=f"{name}_maxpool_{kernel_size}",
                 )(x[0])
             )
 
-        x = layers.Concatenate(name=f"{name}_concat")(x)
+        x = keras.layers.Concatenate(name=f"{name}_concat")(x)
         x = DarknetConvBlock(
             filters,
             kernel_size=1,
@@ -209,21 +208,23 @@ def DarknetConvBlockDepthwise(
     """
 
     if name is None:
-        name = f"conv_block{backend.get_uid('conv_block')}"
+        name = f"conv_block{keras.backend.get_uid('conv_block')}"
 
     model_layers = [
-        layers.DepthwiseConv2D(
+        keras.layers.DepthwiseConv2D(
             kernel_size, strides, padding="same", use_bias=False
         ),
-        layers.BatchNormalization(),
+        keras.layers.BatchNormalization(),
     ]
 
     if activation == "silu":
-        model_layers.append(layers.Lambda(lambda x: keras.activations.swish(x)))
+        model_layers.append(
+            keras.layers.Lambda(lambda x: keras.activations.swish(x))
+        )
     elif activation == "relu":
-        model_layers.append(layers.ReLU())
+        model_layers.append(keras.layers.ReLU())
     elif activation == "leaky_relu":
-        model_layers.append(layers.LeakyReLU(0.1))
+        model_layers.append(keras.layers.LeakyReLU(0.1))
 
     model_layers.append(
         DarknetConvBlock(
@@ -235,7 +236,7 @@ def DarknetConvBlockDepthwise(
 
 
 @keras.utils.register_keras_serializable(package="keras_cv")
-class CrossStagePartial(layers.Layer):
+class CrossStagePartial(keras.layers.Layer):
     """A block used in Cross Stage Partial Darknet.
 
     Args:
@@ -309,8 +310,8 @@ def __init__(
                 )
             )
 
-        self.add = layers.Add()
-        self.concatenate = layers.Concatenate()
+        self.add = keras.layers.Add()
+        self.concatenate = keras.layers.Concatenate()
 
         self.darknet_conv3 = DarknetConvBlock(
             filters, kernel_size=1, strides=1, activation=activation
@@ -360,17 +361,13 @@ def Focus(name=None):
     """  # noqa: E501
 
     def apply(x):
-        return layers.Lambda(
-            lambda x: tf.concat(
-                [
-                    x[..., ::2, ::2, :],
-                    x[..., 1::2, ::2, :],
-                    x[..., ::2, 1::2, :],
-                    x[..., 1::2, 1::2, :],
-                ],
-                axis=-1,
-            ),
-            name=name,
-        )(x)
+        return keras.layers.Concatenate(name=name)(
+            [
+                x[..., ::2, ::2, :],
+                x[..., 1::2, ::2, :],
+                x[..., ::2, 1::2, :],
+                x[..., 1::2, 1::2, :],
+            ],
+        )
 
     return apply