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
Merged

Support PyDataset in Normalization layer adapt methods #21817

Show file tree
Hide file tree
Changes from 1 commit
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
20 changes: 20 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,24 @@ 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.

input_shape = data.element_spec.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)
input_shape = data.element_spec.shape
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
else:
raise NotImplementedError(f"Unsupported data type: {type(data)}")

if not self.built:
self.build(input_shape)
Expand Down
25 changes: 25 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,28 @@ 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)
Loading