EfficentNets no longer train on non-float32 inputs

[hctomkins]

System information

• Have I written custom code (as opposed to using a stock example script provided in TensorFlow): yes
• OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Linux Ubuntu 18.04
• TensorFlow installed from (source or binary): Pip tf-nightly
• TensorFlow version (use command below): 2.10.0-dev20220409
• Python version: 3.7.3
• GPU model and memory: N/A - reproducible on CPU, GPU, and TPU

Describe the problem.
This PR merged to solve an issue with efficientnet normalisation hardcoded a set of float values into the efficientnet architecture.

These magic numbers are fixed to float32 when running in graph mode, and will raise an exception as soon as the network is called on float16 or bfloat16 inputs. bfloat16 inputs are recommended when training on TPUs

Describe the current behavior.
An exception is raised whenever a keras efficientnet is trained on a non-float32 input

Describe the expected behavior.
An exception is not raised whenever a keras efficientnet is trained on a non-float32 input

• Do you want to contribute a PR? (yes/no): yes
• Briefly describe your candidate solution(if contributing): Replace the hardcoded graph constants with:

x = layers.Rescaling(1. / tf.math.sqrt(IMAGENET_STDDEV_RGB))(x)

Standalone code to reproduce the issue.

import tensorflow as tf
from tensorflow.keras.applications import efficientnet as efn

tf.keras.mixed_precision.set_global_policy("mixed_bfloat16")

ds = (
    tf.data.Dataset.from_tensor_slices(
        (
            tf.ones([1, 224, 224, 3], dtype=tf.bfloat16),
            tf.ones([1, 1000], dtype=tf.float32),
        )
    )
    .repeat()
    .batch(batch_size=1)
)

e = efn.EfficientNetB0()
e.compile(loss="binary_crossentropy")
e.fit(ds, steps_per_epoch=1, epochs=1)

[sushreebarsa]

@hctomkins This issue will be closed once the PR is merged.
Thank you!

[google-ml-butler]

Are you satisfied with the resolution of your issue?
Yes
No