Resizing layer with mixed_precision

[robertatdm]

Current Behavior:

The keras_cv.layers.Resizing layer seems to have problems with mixed precision. The following code snippet only works if we don't set the mixed_precision policy.

Expected Behavior:

The following code should work without mixed_precision and with mixed_precision enabled. The snippet works, if it outputs (480, 640, 3).

Steps To Reproduce:

import tensorflow as tf
import keras_cv as kcv
import keras
from keras import mixed_precision
policy = mixed_precision.Policy('mixed_float16')
mixed_precision.set_global_policy(policy)

resizing_layer = kcv.layers.Resizing(
    height=480,
    width=640,
    bounding_box_format='xywh',
    pad_to_aspect_ratio=True,
    name="aug_resizing",
)

def pack_fn(inputs):
    images = tf.cast(tf.random.uniform((480*2, 640*2, 3), minval=0, maxval=255, seed=1234), dtype=tf.uint8)
    classes = tf.constant([1], dtype=tf.int32)
    boxes = tf.constant([[10, 10, 20, 20]])
    return {'images': images, 'bounding_boxes': {'boxes': boxes, 'classes': classes}}

ds = tf.data.Dataset.from_tensor_slices([tf.constant([0])])
ds = ds.map(pack_fn)
ds = ds.map(resizing_layer)

sample = next(iter(ds))
print(sample['images'].shape)

ValueError: Exception encountered when calling layer 'aug_resizing' (type Resizing).

in user code:

    File "/home/robert/.local/lib/python3.10/site-packages/keras_cv/src/layers/preprocessing/base_image_augmentation_layer.py", line 433, in call  *
        outputs = self._format_output(self._augment(inputs), metadata)
    File "/home/robert/.local/lib/python3.10/site-packages/keras_cv/src/layers/preprocessing/resizing.py", line 142, in _augment  *
        outputs = self._batch_augment(inputs)
    File "/home/robert/.local/lib/python3.10/site-packages/keras_cv/src/layers/preprocessing/resizing.py", line 382, in _batch_augment  *
        return self._resize_with_pad(inputs)
    File "/home/robert/.local/lib/python3.10/site-packages/keras_cv/src/layers/preprocessing/resizing.py", line 285, in _resize_with_pad  *
        fn_output_signature=fn_output_signature,

    ValueError: Error in map_fn:
      Expected `fn` to return a:
        RaggedTensorSpec(TensorShape([None, 4]), tf.float16, 1, tf.int64)
      But it returned a:
        RaggedTensorSpec(TensorShape([None, 4]), tf.float32, 1, tf.int64)
        (value=tf.RaggedTensor(values=Tensor("aug_resizing/map/while/RaggedFromTensor/Reshape:0", shape=(None,), dtype=float32), row_splits=Tensor("aug_resizing/map/while/RaggedFromTensor/RaggedFromUniformRowLength/RowPartitionFromUniformRowLength/mul:0", shape=(None,), dtype=int64)))
      To fix, update the `fn_output_signature` (or `dtype`) argument to `map_fn`.

Call arguments received by layer 'aug_resizing' (type Resizing):
  • inputs={'images': 'tf.Tensor(shape=(960, 1280, 3), dtype=uint8)', 'bounding_boxes': {'boxes': 'tf.Tensor(shape=(1, 4), dtype=int32)', 'classes': 'tf.Tensor(shape=(1,), dtype=int32)'}}

Version:

Tensorflow version: 2.15.1 Keras version: 2.15.0 Keras-CV version: 0.8.2

[robertatdm]

Solved the issue, by specifying the dtype argument for the resizing layer to be tf.float32.

resizing_layer = kcv.layers.Resizing(
    height=480,
    width=640,
    bounding_box_format='xywh',
    pad_to_aspect_ratio=True,
    name="aug_resizing",
    dtype=tf.float32    # <-- Manually specify dtype, so the layer further computes in float32
)

Should this bug be closed?

[chunduriv]

@robertatdm,

Sure, Please go ahead and close this issue.