Model registry does not load adapted Normalization-Layer (Keras) correctly

[s_plank]

Hello,

I have a problem with the model registry. I'm not sure if I'm registering the model incorrectly or if it´s a bug.

Here are some code snippets:

import pandas as pd
 
import mlflow
import mlflow.keras
import mlflow.tensorflow
from mlflow.tracking.client import MlflowClient
 
from tensorflow.keras import models, layers, utils, optimizers
 
 
X = pd.DataFrame({'f1': [1,2,3], 'f2': [4,5,6], 'f3': [7,8,9]})
y = pd.DataFrame({'target': [4,5,6]})
 
_layers = []
 
normalize_layer = layers.Normalization(axis=1, name="normalization")
normalize_layer.adapt(X)
_layers.append(normalize_layer)
 
h1_layer = layers.Dense(
                name="h1",
                units=2,
                activation='relu',
            )
_layers.append(h1_layer)
 
out_layer = (
    layers.Dense(name="output", units=y.shape[1], activation="linear")
)
_layers.append(out_layer)
 
denormalize_layer = layers.Normalization(
    axis=1, invert=True, name="denormalization"
)
denormalize_layer.adapt(y)
_layers.append(denormalize_layer)
 
model = models.Sequential(layers=_layers)
 
 
with mlflow.start_run(run_name='keras_bug'):
    mlflow.tensorflow.autolog()
    model.compile(optimizer='adam', loss="mae")
    history = model.fit(
        X,
        y,
        epochs=100,
        shuffle=True,
        validation_split=0.33,
        verbose=1,
    )

I use some random data and build a model with a normalization and denormalization layer. Both layers are adapted to X or y.

Normally these layers store the mean and standard deviation (from the training set) and use those values to normalize the following data.

Afterwards I register the run in the databricks model-registry.

Here are the predictions from the trained - not loaded - model:

model.predict(X)
 
array([[3.9997296],
       [4.1370635],
       [3.6655028]], dtype=float32)

Now I use the run_id and load the logged model:

logged_model = 'runs:/33af506cd02d4a56b525134588c5c30a/model'
loaded_model = mlflow.pyfunc.load_model(logged_model)
 
loaded_model.predict(X)

These are the results:

0
0	-6.247854
1	-6.041854
2	-6.749195

The results are completely different compared to the results from the trained model.

Both predictions are correct if I don´t use the Normalization layers so I think the mean and standard deviations are not stored correctly in the registered model.

I tried to get the mean from the loaded_model but I could not find a way to get these values.

Here is a warning I receive when loading the model:

/databricks/python/lib/python3.9/site-packages/keras/backend.py:451: UserWarning: `tf.keras.backend.set_learning_phase` is deprecated and will be removed after 2020-10-11. To update it, simply pass a True/False value to the `training` argument of the `__call__` method of your layer or model.

I tried to set both layers to trainable=False before the registration:

normalize_layer.trainable=False

but this does not solve the problem.

My environment:

ML-Runtime 12.1

ML-Runtime 12.1 uses tensorflow 2.10. I know that tensorflow 2.10 has a bug in the denormalization layer. This bug is solved with tf2.11 but this is another problem and has nothing to do with this behaviour.

Does anyone have an idea how I could solve this problem?