keras: Bugs found in Normalization() and load_model()

First and foremost I reported these bugs at the tensorflow github section which some of the staff told me to post it here.

Setup: OS: Windows 10 x64 tensorflow-gpu: 2.6.0 Python: 3.9.7 CUDA: 11.2.2 cuDNN: 8.1.1

As in the code below, when loading a model that was normalized by not passing arguments to Normalization() it throws an exception when that model is loaded by load_model(), however before loading the model I can use it without any apparent issues which makes you think it’s all good since Normalization() did NOT complain and took care of the input shape. When loading a model that was normalized by Normalization(input_dim=5) it does NOT thrown any exception since a known shape is specified. That is weird I mean it should warn you that when normalizing it without passing arguments to Normalization() you should expect an exception when loading it. This bug also happens in tensorflow-gpu 2.5.0.

bug_1.py:

import numpy as np
import tensorflow as tf


def main():
    train_data = np.array([[1, 2, 3, 4, 5]])

    train_label = np.array([123])

    normalizer = tf.keras.layers.experimental.preprocessing.Normalization()

    normalizer.adapt(train_data)

    model = tf.keras.Sequential([normalizer, tf.keras.layers.Dense(units=1)])

    model.compile(optimizer=tf.optimizers.Adam(learning_rate=0.1), loss='mean_absolute_error')

    model.fit(train_data, train_label, epochs=3000)

    model.save('test.h5')

    model = tf.keras.models.load_model('test.h5')  # Exception is thrown.

    unseen_data = np.array([[1, 2, 3, 4, 6]])

    prediction = model.predict(unseen_data)

    print(prediction)


if __name__ == "__main__":
    main()

Exception:

Traceback (most recent call last):
  File "E:\Backup\Desktop\bug_1.py", line 28, in <module>
    main()
  File "E:\Backup\Desktop\bug_1.py", line 19, in main
    model = tf.keras.models.load_model('test.h5')
  File "C:\Users\censored\AppData\Local\Programs\Python\Python39\lib\site-packages\keras\saving\save.py", line 200, in load_model
    return hdf5_format.load_model_from_hdf5(filepath, custom_objects,
  File "C:\Users\censored\AppData\Local\Programs\Python\Python39\lib\site-packages\keras\saving\hdf5_format.py", line 180, in load_model_from_hdf5
    model = model_config_lib.model_from_config(model_config,
  File "C:\Users\censored\AppData\Local\Programs\Python\Python39\lib\site-packages\keras\saving\model_config.py", line 52, in model_from_config
    return deserialize(config, custom_objects=custom_objects)
  File "C:\Users\censored\AppData\Local\Programs\Python\Python39\lib\site-packages\keras\layers\serialization.py", line 208, in deserialize
    return generic_utils.deserialize_keras_object(
  File "C:\Users\censored\AppData\Local\Programs\Python\Python39\lib\site-packages\keras\utils\generic_utils.py", line 674, in deserialize_keras_object
    deserialized_obj = cls.from_config(
  File "C:\Users\censored\AppData\Local\Programs\Python\Python39\lib\site-packages\keras\engine\sequential.py", line 434, in from_config
    model.add(layer)
  File "C:\Users\censored\AppData\Local\Programs\Python\Python39\lib\site-packages\tensorflow\python\training\tracking\base.py", line 530, in _method_wrapper
    result = method(self, *args, **kwargs)
  File "C:\Users\censored\AppData\Local\Programs\Python\Python39\lib\site-packages\keras\engine\sequential.py", line 217, in add
    output_tensor = layer(self.outputs[0])
  File "C:\Users\censored\AppData\Local\Programs\Python\Python39\lib\site-packages\keras\engine\base_layer.py", line 976, in __call__
    return self._functional_construction_call(inputs, args, kwargs,
  File "C:\Users\censored\AppData\Local\Programs\Python\Python39\lib\site-packages\keras\engine\base_layer.py", line 1114, in _functional_construction_call
    outputs = self._keras_tensor_symbolic_call(
  File "C:\Users\censored\AppData\Local\Programs\Python\Python39\lib\site-packages\keras\engine\base_layer.py", line 848, in _keras_tensor_symbolic_call
    return self._infer_output_signature(inputs, args, kwargs, input_masks)
  File "C:\Users\censored\AppData\Local\Programs\Python\Python39\lib\site-packages\keras\engine\base_layer.py", line 886, in _infer_output_signature
    self._maybe_build(inputs)
  File "C:\Users\censored\AppData\Local\Programs\Python\Python39\lib\site-packages\keras\engine\base_layer.py", line 2659, in _maybe_build
    self.build(input_shapes)  # pylint:disable=not-callable
  File "C:\Users\censored\AppData\Local\Programs\Python\Python39\lib\site-packages\keras\layers\preprocessing\normalization.py", line 145, in build
    raise ValueError(
ValueError: All `axis` values to be kept must have known shape. Got axis: (-1,), input shape: [None, None], with unknown axis at index: 1

Process finished with exit code 1

In addition to that, there is another bug when you load a model, that model gets a significant loss of accuracy, however that does NOT happen in tensorflow-gpu 2.5.0 which means tensorflow-gpu 2.6.0 has more bugs than expected.

bug_2.py:

import numpy as np
import tensorflow as tf


def main():
    train_data = np.array([[1, 2, 3, 4, 5]])

    train_label = np.array([123])

    normalizer = tf.keras.layers.experimental.preprocessing.Normalization(input_dim=5, axis=None)

    normalizer.adapt(train_data)

    model = tf.keras.Sequential([normalizer, tf.keras.layers.Dense(units=1)])

    model.compile(optimizer=tf.optimizers.Adam(learning_rate=0.1), loss='mean_absolute_error')

    model.fit(train_data, train_label, epochs=3000)

    model.save('test.h5')

    model_loaded = tf.keras.models.load_model('test.h5')

    seen_data = np.array([[1, 2, 3, 4, 5]])

    model_prediction = model.predict(seen_data)

    model_loaded_prediction = model_loaded.predict(seen_data)

    print(f'model prediction:{model_prediction}')  # Expected value close to 123.

    print(f'model loaded prediction: {model_loaded_prediction}')   # Wrong, not even close to 123.


if __name__ == "__main__":
    main()