Repository: keras-team/keras-cv

Issue:
Remove the `training` argument from `.call()` function of all preprocessing layers
According to the discussion with @LukeWood and @fchollet , we will remove the `training` argument from all `.call()` functions of all preprocessing layers in KerasCV.

The reason is to easily build the preprocessing pipeline into a `Sequential` model and run it with `tf.data.Dataset.map()`.
This preprocessing pipeline is separate from the main neural network model, thus do not receive the `training` argument when calling `model.fit()` in the main neural network model. Therefore, managing this `training` argument in the preprocessing pipeline is an extra burden on the users.

After removing the `training` argument, all preprocessing layers are expected only to perform training time behaviors all the time. The user should construct separate pipelines for training and inferencing.

Following are the preprocessing layers in KerasCV that overrides the `call()` function, which needs to be changed in addressing this issue.

`keras_cv/layers/preprocessing/`
- [ ] randomly_zoomed_crop.py
- [ ] mosaic.py
- [ ] base_image_augmentation_layer.py
- [ ] random_crop_and_resize.py
- [ ] vectorized_base_image_augmentation_layer.py
- [ ] resizing.py
- [x] jittered_resize.py

`keras_cv/layers/preprocessing_3d/`
- [ ] group_points_by_bounding_boxes.py
- [ ] base_augmentation_layer_3d.py
- [ ] random_copy_paste.py

Discussion:
soma2000-lang: @james77777778 Please remove my name  ,I am just a contributor, not a maintainer.
LukeWood: Looks great @james77777778 !
LukeWood: (Aside from the change to mosaic)

Patch:
diff --git a/.devcontainer/setup.sh b/.devcontainer/setup.sh
index 6f2daf0..419a861 100644
--- a/.devcontainer/setup.sh
+++ b/.devcontainer/setup.sh
@@ -1,2 +1,4 @@
 sudo pip install --upgrade pip
 sudo pip install -e ".[tests]"
+sudo apt update
+sudo apt install -y clang-format
diff --git a/keras_cv/layers/preprocessing/base_image_augmentation_layer.py b/keras_cv/layers/preprocessing/base_image_augmentation_layer.py
index 536a69a..386fd7d 100644
--- a/keras_cv/layers/preprocessing/base_image_augmentation_layer.py
+++ b/keras_cv/layers/preprocessing/base_image_augmentation_layer.py
@@ -1,4 +1,4 @@
-# Copyright 2022 The KerasCV Authors
+# Copyright 2023 The KerasCV Authors
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -383,25 +383,20 @@ class BaseImageAugmentationLayer(keras.__internal__.layers.BaseRandomLayer):
         """
         return None
 
-    def call(self, inputs, training=True):
+    def call(self, inputs):
         inputs = self._ensure_inputs_are_compute_dtype(inputs)
-        if training:
-            inputs, metadata = self._format_inputs(inputs)
-            images = inputs[IMAGES]
-            if images.shape.rank == 3:
-                return self._format_output(self._augment(inputs), metadata)
-            elif images.shape.rank == 4:
-                return self._format_output(
-                    self._batch_augment(inputs), metadata
-                )
-            else:
-                raise ValueError(
-                    "Image augmentation layers are expecting inputs to be "
-                    "rank 3 (HWC) or 4D (NHWC) tensors. Got shape: "
-                    f"{images.shape}"
-                )
+        inputs, metadata = self._format_inputs(inputs)
+        images = inputs[IMAGES]
+        if images.shape.rank == 3:
+            return self._format_output(self._augment(inputs), metadata)
+        elif images.shape.rank == 4:
+            return self._format_output(self._batch_augment(inputs), metadata)
         else:
-            return inputs
+            raise ValueError(
+                "Image augmentation layers are expecting inputs to be "
+                "rank 3 (HWC) or 4D (NHWC) tensors. Got shape: "
+                f"{images.shape}"
+            )
 
     def _augment(self, inputs):
         raw_image = inputs.get(IMAGES, None)
diff --git a/keras_cv/layers/preprocessing/mosaic.py b/keras_cv/layers/preprocessing/mosaic.py
index d6fadf3..22a3e68 100644
--- a/keras_cv/layers/preprocessing/mosaic.py
+++ b/keras_cv/layers/preprocessing/mosaic.py
@@ -269,17 +269,16 @@ class Mosaic(VectorizedBaseImageAugmentationLayer):
         self._validate_inputs(inputs)
         return super()._batch_augment(inputs)
 
-    def call(self, inputs, training=True):
-        if training is True:
-            _, metadata = self._format_inputs(inputs)
-            if metadata[BATCHED] is not True:
-                raise ValueError(
-                    "Mosaic received a single image to `call`. The "
-                    "layer relies on combining multiple examples, and as such "
-                    "will not behave as expected. Please call the layer with 4 "
-                    "or more samples."
-                )
-        return super().call(inputs=inputs, training=training)
+    def call(self, inputs):
+        _, metadata = self._format_inputs(inputs)
+        if metadata[BATCHED] is not True:
+            raise ValueError(
+                "Mosaic received a single image to `call`. The "
+                "layer relies on combining multiple examples, and as such "
+                "will not behave as expected. Please call the layer with 4 "
+                "or more samples."
+            )
+        return super().call(inputs=inputs)
 
     def _validate_inputs(self, inputs):
         images = inputs.get(IMAGES, None)
diff --git a/keras_cv/layers/preprocessing/random_crop_and_resize.py b/keras_cv/layers/preprocessing/random_crop_and_resize.py
index 0f5c278..5af4cc7 100644
--- a/keras_cv/layers/preprocessing/random_crop_and_resize.py
+++ b/keras_cv/layers/preprocessing/random_crop_and_resize.py
@@ -128,24 +128,6 @@ class RandomCropAndResize(BaseImageAugmentationLayer):
 
         return [[y1, x1, y2, x2]]
 
-    def call(self, inputs, training=True):
-        if training:
-            return super().call(inputs, training)
-        else:
-            inputs = self._ensure_inputs_are_compute_dtype(inputs)
-            inputs, meta_data = self._format_inputs(inputs)
-            output = inputs
-            # self._resize() returns valid results for both batched and
-            # unbatched
-            output["images"] = self._resize(inputs["images"])
-
-            if "segmentation_masks" in inputs:
-                output["segmentation_masks"] = self._resize(
-                    inputs["segmentation_masks"], interpolation="nearest"
-                )
-
-            return self._format_output(output, meta_data)
-
     def compute_image_signature(self, images):
         return tf.TensorSpec(
             shape=(self.target_size[0], self.target_size[1], images.shape[-1]),
diff --git a/keras_cv/layers/preprocessing/randomly_zoomed_crop.py b/keras_cv/layers/preprocessing/randomly_zoomed_crop.py
index b3c79c1..3e736d0 100644
--- a/keras_cv/layers/preprocessing/randomly_zoomed_crop.py
+++ b/keras_cv/layers/preprocessing/randomly_zoomed_crop.py
@@ -17,9 +17,6 @@ from keras import backend
 from tensorflow import keras
 
 from keras_cv import core
-from keras_cv.layers.preprocessing.vectorized_base_image_augmentation_layer import (
-    IMAGES,
-)
 from keras_cv.layers.preprocessing.vectorized_base_image_augmentation_layer import (
     VectorizedBaseImageAugmentationLayer,
 )
@@ -259,19 +256,6 @@ class RandomlyZoomedCrop(VectorizedBaseImageAugmentationLayer):
                 axis=1,
             )
 
-    def call(self, inputs, training=True):
-        if training:
-            return super().call(inputs, training)
-        else:
-            inputs = self._ensure_inputs_are_compute_dtype(inputs)
-            inputs, meta_data = self._format_inputs(inputs)
-            output = inputs
-            # self._resize() returns valid results for both batched and
-            # unbatched
-            output[IMAGES] = self._resize(inputs[IMAGES])
-
-            return self._format_output(output, meta_data)
-
     def _resize(self, images, **kwargs):
         resizing_layer = keras.layers.Resizing(
             self.height, self.width, **kwargs
diff --git a/keras_cv/layers/preprocessing/resizing.py b/keras_cv/layers/preprocessing/resizing.py
index 030fe22..37b3372 100644
--- a/keras_cv/layers/preprocessing/resizing.py
+++ b/keras_cv/layers/preprocessing/resizing.py
@@ -273,22 +273,6 @@ class Resizing(BaseImageAugmentationLayer):
         inputs["images"] = images
         return inputs
 
-    def call(self, inputs, training=True):
-        inputs = self._ensure_inputs_are_compute_dtype(inputs)
-        inputs, metadata = self._format_inputs(inputs)
-        self._check_inputs(inputs)
-        images = inputs["images"]
-        if images.shape.rank == 3:
-            return self._format_output(self._augment(inputs), metadata)
-        elif images.shape.rank == 4:
-            return self._format_output(self._batch_augment(inputs), metadata)
-        else:
-            raise ValueError(
-                "Image augmentation layers are expecting inputs to be "
-                "rank 3 (HWC) or 4D (NHWC) tensors. Got shape: "
-                f"{images.shape}"
-            )
-
     def _check_inputs(self, inputs):
         for key in inputs:
             if key not in supported_keys:
diff --git a/keras_cv/layers/preprocessing/vectorized_base_image_augmentation_layer.py b/keras_cv/layers/preprocessing/vectorized_base_image_augmentation_layer.py
index 5e642e2..1573a21 100644
--- a/keras_cv/layers/preprocessing/vectorized_base_image_augmentation_layer.py
+++ b/keras_cv/layers/preprocessing/vectorized_base_image_augmentation_layer.py
@@ -348,24 +348,18 @@ class VectorizedBaseImageAugmentationLayer(
             result[key] = inputs[key]
         return result
 
-    def call(self, inputs, training=True):
-        # TODO(lukewood): remove training=False behavior.
+    def call(self, inputs):
         inputs = self._ensure_inputs_are_compute_dtype(inputs)
-        if training:
-            inputs, metadata = self._format_inputs(inputs)
-            images = inputs[IMAGES]
-            if images.shape.rank == 3 or images.shape.rank == 4:
-                return self._format_output(
-                    self._batch_augment(inputs), metadata
-                )
-            else:
-                raise ValueError(
-                    "Image augmentation layers are expecting inputs to be "
-                    "rank 3 (HWC) or 4D (NHWC) tensors. Got shape: "
-                    f"{images.shape}"
-                )
+        inputs, metadata = self._format_inputs(inputs)
+        images = inputs[IMAGES]
+        if images.shape.rank == 3 or images.shape.rank == 4:
+            return self._format_output(self._batch_augment(inputs), metadata)
         else:
-            return inputs
+            raise ValueError(
+                "Image augmentation layers are expecting inputs to be "
+                "rank 3 (HWC) or 4D (NHWC) tensors. Got shape: "
+                f"{images.shape}"
+            )
 
     def _format_inputs(self, inputs):
         metadata = {IS_DICT: True, USE_TARGETS: False}
diff --git a/keras_cv/losses/__init__.py b/keras_cv/losses/__init__.py
index a6f9429..2f575da 100644
--- a/keras_cv/losses/__init__.py
+++ b/keras_cv/losses/__init__.py
@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from keras_cv.losses.center_net_box_loss import CenterNetBoxLoss
 from keras_cv.losses.focal import FocalLoss
 from keras_cv.losses.giou_loss import GIoULoss
 from keras_cv.losses.iou_loss import IoULoss
diff --git a/keras_cv/losses/center_net_box_loss.py b/keras_cv/losses/center_net_box_loss.py
new file mode 100644
index 0000000..d48e7f9
--- /dev/null
+++ b/keras_cv/losses/center_net_box_loss.py
@@ -0,0 +1,133 @@
+# Copyright 2023 The KerasCV Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import math
+
+import tensorflow as tf
+from tensorflow import keras
+
+from keras_cv.bounding_box_3d import CENTER_XYZ_DXDYDZ_PHI
+
+
+def l1(y_true, y_pred, sigma=9.0):
+    """Computes element-wise l1 loss."""
+
+    absolute_difference = tf.abs(y_pred - y_true)
+    loss = tf.where(
+        absolute_difference < 1.0 / sigma,
+        0.5 * sigma * absolute_difference**2,
+        absolute_difference - 0.5 / sigma,
+    )
+
+    return tf.reduce_sum(loss, axis=-1)
+
+
+class CenterNetBoxLoss(keras.losses.Loss):
+    """Implements a bin-based box regression loss for 3D bounding boxes.
+
+    This loss is meant for use as a box loss for
+    `keras_cv.models.MultiHeadCenterPillar`.
+
+    Reference: https://arxiv.org/abs/1812.04244
+
+    CenterNetBoxLoss uses L1 loss on the individual components of
+    boxes, with the exception of the bin-based heading component of each box,
+    where the bin indicator outputs use cross entropy loss, and the heading
+    residual uses L1 loss. The position (x/y/z) components of predictions are
+    absolute, while the size components are normalized to the anchor size.
+
+    Ground truth boxes are expected to follow the CENTER_XYZ_DXDYDZ_PHI format.
+    Refer to https://github.com/keras-team/keras-cv/blob/master/keras_cv/bounding_box_3d/formats.py
+    for more details on supported bounding box formats.
+
+    Box predictions are expected to be in CenterPillar heatmap-encoded format.
+
+    Args:
+        num_heading_bins: int, number of bins used for predicting box heading.
+        anchor_size: list of 3 ints, anchor sizes for the x, y, and z axes.
+    """
+
+    def __init__(self, num_heading_bins, anchor_size, **kwargs):
+        super().__init__(**kwargs)
+        self.num_heading_bins = num_heading_bins
+        self.anchor_size = anchor_size
+
+    def heading_regression_loss(self, heading_true, heading_pred):
+        # Set the heading to within 0 -> 2pi
+        heading_true = tf.math.floormod(heading_true, 2 * math.pi)
+
+        # Divide 2pi into bins. shifted by 0.5 * angle_per_class.
+        angle_per_class = (2 * math.pi) / self.num_heading_bins
+        shift_angle = tf.math.floormod(
+            heading_true + angle_per_class / 2, 2 * math.pi
+        )
+
+        heading_bin_label_float = tf.math.floordiv(shift_angle, angle_per_class)
+        heading_bin_label = tf.cast(heading_bin_label_float, dtype=tf.int32)
+        heading_res_label = shift_angle - (
+            heading_bin_label_float * angle_per_class + angle_per_class / 2.0
+        )
+        heading_res_norm_label = heading_res_label / (angle_per_class / 2.0)
+
+        heading_bin_one_hot = tf.one_hot(
+            heading_bin_label, self.num_heading_bins
+        )
+        loss_heading_bin = tf.nn.softmax_cross_entropy_with_logits(
+            labels=heading_bin_one_hot,
+            logits=heading_pred[..., : self.num_heading_bins],
+        )
+        loss_heading_res = l1(
+            tf.reduce_sum(
+                heading_pred[..., self.num_heading_bins :]
+                * heading_bin_one_hot,
+                axis=-1,
+                keepdims=True,
+            ),
+            heading_res_norm_label[..., tf.newaxis],
+        )
+
+        return loss_heading_bin + loss_heading_res
+
+    def regression_loss(self, y_true, y_pred):
+        position_loss = l1(y_true[:, :3], y_pred[:, :3])
+
+        heading_loss = self.heading_regression_loss(
+            y_true[:, CENTER_XYZ_DXDYDZ_PHI.PHI], y_pred[:, 3:-3]
+        )
+
+        # Size loss
+        size_norm_label = y_true[:, 3:6] / self.anchor_size
+        size_norm_pred = y_pred[:, -3:] + 1.0
+        size_loss = l1(size_norm_pred, size_norm_label)
+
+        # TODO(ianstenbit): Add IoU3D Loss.
+
+        return position_loss + heading_loss + size_loss
+
+    def call(self, y_true, y_pred):
+        return tf.map_fn(
+            lambda y_true_and_pred: self.regression_loss(
+                y_true_and_pred[0], y_true_and_pred[1]
+            ),
+            (y_true, y_pred),
+            tf.TensorSpec((y_pred.shape[1])),
+        )
+
+    def get_config(self):
+        config = {
+            "num_heading_bins": self.num_heading_bins,
+            "anchor_size": self.anchor_size,
+        }
+        base_config = super().get_config()
+        return dict(list(base_config.items()) + list(config.items()))
diff --git a/keras_cv/models/__init__.py b/keras_cv/models/__init__.py
index cc49544..f0acb3b 100644
--- a/keras_cv/models/__init__.py
+++ b/keras_cv/models/__init__.py
@@ -27,6 +27,9 @@ from keras_cv.models.backbones.resnet_v1.resnet_v1_backbone import (
 from keras_cv.models.backbones.resnet_v1.resnet_v1_backbone import (
     ResNet152Backbone,
 )
+from keras_cv.models.backbones.resnet_v1.resnet_v1_backbone import (
+    ResNetBackbone,
+)
 from keras_cv.models.backbones.resnet_v2.resnet_v2_backbone import (
     ResNet18V2Backbone,
 )


Test patch:
diff --git a/keras_cv/layers/preprocessing/mosaic_test.py b/keras_cv/layers/preprocessing/mosaic_test.py
index 49980eb..659ed1a 100644
--- a/keras_cv/layers/preprocessing/mosaic_test.py
+++ b/keras_cv/layers/preprocessing/mosaic_test.py
@@ -1,4 +1,4 @@
-# Copyright 2022 The KerasCV Authors
+# Copyright 2023 The KerasCV Authors
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
diff --git a/keras_cv/layers/preprocessing/random_aspect_ratio_test.py b/keras_cv/layers/preprocessing/random_aspect_ratio_test.py
index 3ad4e26..de86f55 100644
--- a/keras_cv/layers/preprocessing/random_aspect_ratio_test.py
+++ b/keras_cv/layers/preprocessing/random_aspect_ratio_test.py
@@ -27,15 +27,6 @@ class RandomAspectRatioTest(tf.test.TestCase):
         output = layer(image, training=True)
         self.assertNotEqual(output.shape, image.shape)
 
-    def test_inference_preserves_image(self):
-        # Checks if original and augmented images are different
-        input_image_shape = (8, 100, 100, 3)
-        image = tf.random.uniform(shape=input_image_shape)
-
-        layer = layers.RandomAspectRatio(factor=(0.9, 1.1))
-        output = layer(image, training=False)
-        self.assertAllClose(image, output)
-
     def test_grayscale(self):
         # Checks if original and augmented images are different
         input_image_shape = (8, 100, 100, 1)
diff --git a/keras_cv/layers/preprocessing/random_channel_shift_test.py b/keras_cv/layers/preprocessing/random_channel_shift_test.py
index f92bdb2..bb22951 100644
--- a/keras_cv/layers/preprocessing/random_channel_shift_test.py
+++ b/keras_cv/layers/preprocessing/random_channel_shift_test.py
@@ -12,7 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import numpy as np
 import tensorflow as tf
 from absl.testing import parameterized
 
@@ -115,11 +114,3 @@ class RandomChannelShiftTest(tf.test.TestCase, parameterized.TestCase):
         self.assertEqual(reconstructed_layer.value_range, layer.value_range)
         self.assertEqual(reconstructed_layer.seed, layer.seed)
         self.assertEqual(reconstructed_layer.channels, layer.channels)
-
-    def test_inference(self):
-        layer = preprocessing.RandomChannelShift(
-            factor=0.8, value_range=(0, 255)
-        )
-        inputs = np.random.randint(0, 255, size=(224, 224, 3))
-        output = layer(inputs, training=False)
-        self.assertAllClose(inputs, output)
diff --git a/keras_cv/layers/preprocessing/random_color_jitter_test.py b/keras_cv/layers/preprocessing/random_color_jitter_test.py
index 4405189..3831f2a 100644
--- a/keras_cv/layers/preprocessing/random_color_jitter_test.py
+++ b/keras_cv/layers/preprocessing/random_color_jitter_test.py
@@ -12,7 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import numpy as np
 import tensorflow as tf
 from absl.testing import parameterized
 
@@ -104,16 +103,3 @@ class RandomColorJitterTest(tf.test.TestCase, parameterized.TestCase):
             reconstructed_layer.saturation_factor, layer.saturation_factor
         )
         self.assertEqual(reconstructed_layer.hue_factor, layer.hue_factor)
-
-    # Test 5: Check if inference model is OK.
-    def test_inference(self):
-        layer = preprocessing.RandomColorJitter(
-            value_range=(0, 255),
-            brightness_factor=0.5,
-            contrast_factor=(0.5, 0.9),
-            saturation_factor=(0.5, 0.9),
-            hue_factor=0.5,
-        )
-        inputs = np.random.randint(0, 255, size=(224, 224, 3))
-        output = layer(inputs, training=False)
-        self.assertAllClose(inputs, output)
diff --git a/keras_cv/layers/preprocessing/random_crop_and_resize_test.py b/keras_cv/layers/preprocessing/random_crop_and_resize_test.py
index a5dacf9..ca32a8c 100644
--- a/keras_cv/layers/preprocessing/random_crop_and_resize_test.py
+++ b/keras_cv/layers/preprocessing/random_crop_and_resize_test.py
@@ -59,21 +59,6 @@ class RandomCropAndResizeTest(tf.test.TestCase, parameterized.TestCase):
         self.assertAllEqual(output.shape, (4, 224, 224, 1))
         self.assertNotAllClose(output, input_image_resized)
 
-    def test_preserves_image(self):
-        image_shape = (self.batch_size, self.height, self.width, 3)
-        image = tf.random.uniform(shape=image_shape)
-
-        layer = preprocessing.RandomCropAndResize(
-            target_size=self.target_size,
-            aspect_ratio_factor=(3 / 4, 4 / 3),
-            crop_area_factor=(0.8, 1.0),
-        )
-
-        input_resized = tf.image.resize(image, self.target_size)
-        output = layer(image, training=False)
-
-        self.assertAllClose(output, input_resized)
-
     @parameterized.named_parameters(
         ("Not tuple or list", dict()),
         ("Length not equal to 2", [1, 2, 3]),
diff --git a/keras_cv/layers/preprocessing/random_flip_test.py b/keras_cv/layers/preprocessing/random_flip_test.py
index e92de89..f4933cc 100644
--- a/keras_cv/layers/preprocessing/random_flip_test.py
+++ b/keras_cv/layers/preprocessing/random_flip_test.py
@@ -77,13 +77,6 @@ class RandomFlipTest(tf.test.TestCase, parameterized.TestCase):
             actual_output = layer(inp, training=True)
             self.assertAllClose(expected_output, actual_output)
 
-    def test_random_flip_inference(self):
-        input_images = np.random.random((2, 5, 8, 3)).astype(np.float32)
-        expected_output = input_images
-        layer = RandomFlip()
-        actual_output = layer(input_images, training=False)
-        self.assertAllClose(expected_output, actual_output)
-
     def test_random_flip_default(self):
         input_images = np.random.random((2, 5, 8, 3)).astype(np.float32)
         expected_output = np.flip(input_images, axis=2)
diff --git a/keras_cv/layers/preprocessing/random_rotation_test.py b/keras_cv/layers/preprocessing/random_rotation_test.py
index d5dbe06..85c9020 100644
--- a/keras_cv/layers/preprocessing/random_rotation_test.py
+++ b/keras_cv/layers/preprocessing/random_rotation_test.py
@@ -26,13 +26,6 @@ class RandomRotationTest(tf.test.TestCase):
         actual_output = layer(input_images, training=True)
         self.assertEqual(expected_output.shape, actual_output.shape)
 
-    def test_random_rotation_inference(self):
-        input_images = np.random.random((2, 5, 8, 3)).astype(np.float32)
-        expected_output = input_images
-        layer = RandomRotation(0.5)
-        actual_output = layer(input_images, training=False)
-        self.assertAllClose(expected_output, actual_output)
-
     def test_random_rotation_on_batched_images_independently(self):
         image = tf.random.uniform((100, 100, 3))
         batched_images = tf.stack((image, image), axis=0)
diff --git a/keras_cv/layers/preprocessing/random_translation_test.py b/keras_cv/layers/preprocessing/random_translation_test.py
index ad1ea75..6b502e7 100644
--- a/keras_cv/layers/preprocessing/random_translation_test.py
+++ b/keras_cv/layers/preprocessing/random_translation_test.py
@@ -185,13 +185,6 @@ class RandomTranslationTest(tf.test.TestCase, parameterized.TestCase):
             expected_output = np.reshape(expected_output, (1, 5, 5, 1))
             self.assertAllEqual(expected_output, output_image)
 
-    def test_random_translation_inference(self):
-        input_images = np.random.random((2, 5, 8, 3)).astype(np.float32)
-        expected_output = input_images
-        layer = preprocessing.RandomTranslation(0.5, 0.5)
-        actual_output = layer(input_images, training=False)
-        self.assertAllClose(expected_output, actual_output)
-
     def test_random_translation_on_batched_images_independently(self):
         image = tf.random.uniform(shape=(100, 100, 3))
         input_images = tf.stack([image, image], axis=0)
diff --git a/keras_cv/layers/preprocessing/random_zoom_test.py b/keras_cv/layers/preprocessing/random_zoom_test.py
index b967ef2..d266245 100644
--- a/keras_cv/layers/preprocessing/random_zoom_test.py
+++ b/keras_cv/layers/preprocessing/random_zoom_test.py
@@ -109,13 +109,6 @@ class RandomZoomTest(tf.test.TestCase, parameterized.TestCase):
             expected_output = np.reshape(expected_output, (1, 5, 5, 1))
             self.assertAllEqual(expected_output, output_image)
 
-    def test_random_zoom_inference(self):
-        input_images = np.random.random((2, 5, 8, 3)).astype(np.float32)
-        expected_output = input_images
-        layer = RandomZoom(0.5, 0.5)
-        actual_output = layer(input_images, training=False)
-        self.assertAllClose(expected_output, actual_output)
-
     def test_random_zoom_on_batched_images_independently(self):
         image = tf.random.uniform(shape=(100, 100, 3))
         input_images = tf.stack([image, image], axis=0)
diff --git a/keras_cv/layers/preprocessing/randomly_zoomed_crop_test.py b/keras_cv/layers/preprocessing/randomly_zoomed_crop_test.py
index 2f21714..baf2c31 100644
--- a/keras_cv/layers/preprocessing/randomly_zoomed_crop_test.py
+++ b/keras_cv/layers/preprocessing/randomly_zoomed_crop_test.py
@@ -61,22 +61,6 @@ class RandomlyZoomedCropTest(tf.test.TestCase, parameterized.TestCase):
         self.assertAllEqual(output.shape, (4, 224, 224, 1))
         self.assertNotAllClose(output, input_image_resized)
 
-    def test_preserves_image(self):
-        image_shape = (self.batch_size, self.height, self.width, 3)
-        image = tf.random.uniform(shape=image_shape)
-
-        layer = preprocessing.RandomlyZoomedCrop(
-            height=self.target_size[0],
-            width=self.target_size[1],
-            aspect_ratio_factor=(3 / 4, 4 / 3),
-            zoom_factor=(0.8, 1.0),
-        )
-
-        input_resized = tf.image.resize(image, self.target_size)
-        output = layer(image, training=False)
-
-        self.assertAllClose(output, input_resized)
-
     @parameterized.named_parameters(
         ("Not tuple or list", dict()),
         ("Length not equal to 2", [1, 2, 3]),
diff --git a/keras_cv/layers/preprocessing/resizing_test.py b/keras_cv/layers/preprocessing/resizing_test.py
index 9656b59..80f4dc7 100644
--- a/keras_cv/layers/preprocessing/resizing_test.py
+++ b/keras_cv/layers/preprocessing/resizing_test.py
@@ -176,17 +176,6 @@ class ResizingTest(tf.test.TestCase, parameterized.TestCase):
         self.assertNotIsInstance(outputs, tf.RaggedTensor)
         self.assertAllEqual(expected_output, outputs)
 
-    def test_raises_with_segmap(self):
-        inputs = {
-            "images": np.array([[[1], [2]], [[3], [4]]], dtype="float64"),
-            "segmentation_map": np.array(
-                [[[1], [2]], [[3], [4]]], dtype="float64"
-            ),
-        }
-        layer = cv_layers.Resizing(2, 2)
-        with self.assertRaises(ValueError):
-            layer(inputs)
-
     def test_output_dtypes(self):
         inputs = np.array([[[1], [2]], [[3], [4]]], dtype="float64")
         layer = cv_layers.Resizing(2, 2)
diff --git a/keras_cv/losses/center_net_box_loss_test.py b/keras_cv/losses/center_net_box_loss_test.py
new file mode 100644
index 0000000..5fe6139
--- /dev/null
+++ b/keras_cv/losses/center_net_box_loss_test.py
@@ -0,0 +1,43 @@
+# Copyright 2023 The KerasCV Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import tensorflow as tf
+from absl.testing import parameterized
+
+import keras_cv
+
+
+class CenterNetBoxLoss(tf.test.TestCase, parameterized.TestCase):
+    @parameterized.named_parameters(
+        (
+            "none",
+            "none",
+            (
+                2,
+                10,
+            ),
+        ),
+        ("sum", "sum", ()),
+        ("sum_over_batch_size", "sum_over_batch_size", ()),
+    )
+    def test_proper_output_shapes(self, reduction, target_size):
+        loss = keras_cv.losses.CenterNetBoxLoss(
+            num_heading_bins=4, anchor_size=[1.0, 1.0, 1.0], reduction=reduction
+        )
+        result = loss(
+            y_true=tf.random.uniform((2, 10, 7)),
+            # Predictions have xyz,lwh, and 2*4 values for heading.
+            y_pred=tf.random.uniform((2, 10, 6 + 2 * 4)),
+        )
+        self.assertEqual(result.shape, target_size)