Skip to content

Support PyDataset in Normalization layer adapt methods #21817

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 8 commits into from
Nov 27, 2025
+46 −1
Merged

Support PyDataset in Normalization layer adapt methods #21817

Show file tree
Hide file tree
Changes from 3 commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
6a9aa5f
support pydataset in `adapt` for norm layers
danielenricocahall Nov 3, 2025
944082f
address gemini comments around type error and duplication
danielenricocahall Nov 3, 2025
ac460d9
add some slightly more robust checks
danielenricocahall Nov 3, 2025
87db5b3
simplify logic for pydataset support
danielenricocahall Nov 9, 2025
c279453
`Merge branch 'master' into support-pydataset-in-normalization-adapt
danielenricocahall Nov 11, 2025
15d1562
Merge branch 'keras-team:master' into support-pydataset-in-normalizat…
danielenricocahall Nov 23, 2025
0ab5a71
Merge branch 'master' into support-pydataset-in-normalization-adapt
danielenricocahall Nov 25, 2025
883822c
updated based on feedback
danielenricocahall Nov 25, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
22 changes: 22 additions & 0 deletions keras/src/layers/preprocessing/normalization.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,9 @@
from keras.src import ops
from keras.src.api_export import keras_export
from keras.src.layers.preprocessing.data_layer import DataLayer
from keras.src.trainers.data_adapters import get_data_adapter
from keras.src.utils.module_utils import tensorflow as tf
from keras.utils import PyDataset


@keras_export("keras.layers.Normalization")
Expand Down Expand Up @@ -229,6 +231,26 @@ def adapt(self, data):
# Batch dataset if it isn't batched
data = data.batch(128)
input_shape = tuple(data.element_spec.shape)
elif isinstance(data, PyDataset):
# as PyDatasets returns tuples of input/annotation pairs
adapter = get_data_adapter(data)
tf_dataset = adapter.get_tf_dataset()
if len(tf_dataset.element_spec) == 1:
# just x
data = tf_dataset.map(lambda x: x)
Copy link
Collaborator

@fchollet fchollet Nov 5, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This isn't doing anything, and can be removed

Copy link
Contributor Author

@danielenricocahall danielenricocahall Nov 8, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I did that because:

tf_dataset.element_spec
(TensorSpec(shape=(None, 32, 3), dtype=tf.float64, name=None),)
tf_dataset.map(lambda x: x).element_spec
TensorSpec(shape=(None, 32, 3), dtype=tf.float64, name=None)

otherwise, there's a E AttributeError: 'tuple' object has no attribute 'shape' since the former returns a tuple. Alternatively, without the map we could get the shape via tf_dataset.element_spec[0].shape.

elif len(tf_dataset.element_spec) == 2:
Copy link
Collaborator

@fchollet fchollet Nov 5, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It seems fine to only support PyDataset instances that yield data samples (no targets) since that is what we expect for tf.data.Dataset instances. No?

Copy link
Contributor Author

@danielenricocahall danielenricocahall Nov 8, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Possibly - I was thinking there may be users who want to provide a PyDataset with samples and targets and compute the mean/variance from the samples, but it may not be necessary. I could remove it if we don't think that's a use-case

Copy link
Contributor Author

@danielenricocahall danielenricocahall Nov 9, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Updated to use args instead so we don't have these baked in conditionals but will support these variations - what do you think?

# (x, y) pairs
data = tf_dataset.map(lambda x, y: x)
elif len(tf_dataset.element_spec) == 3:
# (x, y, sample_weight) tuples
data = tf_dataset.map(lambda x, y, z: x)
input_shape = data.element_spec.shape
Copy link

@limzikiki limzikiki Nov 3, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Coming from your comment.
What I did in my solution is:

        elif isinstance(data, keras.utils.PyDataset):
            sample_input = data[0][0] # pydataset should return a tuple with first element being the data
            if isinstance(sample_input, np.ndarray) or backend.is_tensor(sample_input):
                input_shape = sample_input.shape
            else:
                raise ValueError(f"Unsupported data type: {type(sample_input)} returned from the PyDataset")

The advantage of my option lies in the fact that we don’t need to perform excessive transformations to tf tensors just for the sake of size estimation. PyDataset is also used for experimentation, and when the dataset is too large to be read into RAM, which is common for workstations and personal devices, transforming PyDataset into a TF Tensor will fail due to memmory allocation. However on contrary the drawback of my solution is that it retrieves the first batch, and during that the first batch might get changed and second retrieval of the first batch might return not the same output (If someone implemented non-idempotent PyDataset, but i think then it is a user problem then). Also why retrieval of the first batch is feasable solution is because shape of all the elements across all the batches must be identical for normalization to work correctly.

Considering strategic direction of Keras to move away from being solely dependant on TensorFlow, adding transformation to tensorflow creates a technical debt that Keras team should later take care off.

I am open for a discussion

Copy link
Contributor Author

@danielenricocahall danielenricocahall Nov 3, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thank you for the thorough response! Yes, I will defer to the core developers' judgement for this. Happy to revise and infer the shape based on sampling a batch if we think that's the better approach.

else:
raise TypeError(
f"Unsupported data type: {type(data)}. `adapt` supports "
f"`np.ndarray`, backend tensors, `tf.data.Dataset`, and "
f"`keras.utils.PyDataset`."
)

if not self.built:
self.build(input_shape)
Expand Down
32 changes: 32 additions & 0 deletions keras/src/layers/preprocessing/normalization_test.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -169,3 +169,35 @@ def test_normalization_with_scalar_mean_var(self):
input_data = np.array([[1, 2, 3]], dtype="float32")
layer = layers.Normalization(mean=3.0, variance=2.0)
layer(input_data)

@parameterized.parameters([("x",), ("x_and_y",), ("x_y_and_weights")])
def test_adapt_pydataset_compat(self, pydataset_type):
import keras

class CustomDataset(keras.utils.PyDataset):
def __len__(self):
return 100

def __getitem__(self, idx):
x = np.random.rand(32, 32, 3)
y = np.random.randint(0, 10, size=(1,))
weights = np.random.randint(0, 10, size=(1,))
if pydataset_type == "x":
return x
elif pydataset_type == "x_and_y":
return x, y
elif pydataset_type == "x_y_and_weights":
return x, y, weights
else:
raise NotImplementedError(pydataset_type)

normalizer = keras.layers.Normalization()
normalizer.adapt(CustomDataset())
self.assertTrue(normalizer.built)
self.assertIsNotNone(normalizer.mean)
self.assertIsNotNone(normalizer.variance)
self.assertEqual(normalizer.mean.shape[-1], 3)
self.assertEqual(normalizer.variance.shape[-1], 3)
sample_input = np.random.rand(1, 32, 32, 3)
output = normalizer(sample_input)
self.assertEqual(output.shape, (1, 32, 32, 3))