keras.layers.Normalization#adapt() does not support keras.utils.PyDataset as input.

[limzikiki]

Expected behavior

When I create a custom PyDataset and then pass it into the Normalization adapt method, the mean and variance should be computed successfully, as for a TensorFlow dataset.

Actual behavior

Method adapt does not handle the PyDataset like inputs.

Traceback (most recent call last):
  File "/home/user/project/keras_bug_sample.py", line 14, in <module>
    normalizer.adapt(CustomDataset)
  File "/home/user/.pyenv/versions/3.12.8/lib/python3.12/site-packages/keras/src/layers/preprocessing/normalization.py", line 230, in adapt
    self.build(input_shape)
               ^^^^^^^^^^^
UnboundLocalError: cannot access local variable 'input_shape' where it is not associated with a value

Steps to reproduce

import keras

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

    def __getitem__(self, idx):
        # Generate dummy data
        x = np.random.rand(32, 32, 3)
        y = np.random.randint(0, 10, size=(1,))
        return x, y

normalizer = keras.layers.Normalization()
normalizer.adapt(CustomDataset)

[sonali-kumari1]

Hi @limzikiki -
Thanks for reporting this. The error UnboundLocalError: cannot access local variable 'input_shape' where it is not associated with a value occurs because adapt() does not directly support custom datasets like keras.utils.PyDataset. The adapt() method expects data to be in the form of tf.data.Dataset, NumPy array, or a backend-native eager tensor. Since you are passing keras.utils.PyDataset, adapt() is not able to extract the input_shape which leads to the UnboundLocalError. To resolve this you can convert your custom dataset into tf.data.Dataset and ensure that dataset is properly batched to compute the mean and variance correctly.

[limzikiki]

I ended up fixing it by subclassing the Normalization class and introducing support for the PyDataset type inputs. I was wondering if the Keras community would also benefit from this feature. Basically, what I did in my solution is that I fetched the first batch, and then, based on its shape, I built the normalisation. I don't know if that's a proper solution. If that’s okay, I can open a PR with the changes.

[VarunS1997]

Please go ahead and make a PR, @limzikiki -- thank you!