Saving normalization values in Keras model

I have a Keras model for which I would like to save the normalization values in the model object itself for easier portability.

I'm using sklearn's StandardScaler() to normalize my data, so I simply want to save the mean_ and var_ attributes from the scaler to the model, save the model, and when I reload the model have access to these attributes.

Currently when I reload the model the attributes I added are not there. What is the correct way of doing this ?

Code:

# Normalize data
scaler = StandardScaler()
scaler.fit(X_train)
...

# Create model
model = Sequential(...)

# Compile and train
...

# Save model with normalization mean and var
model.normalization_mean = scaler.mean_
model.normalization_var  = scaler.var_

keras.models.save_model(model = model, 
                        filepath = ...)

# Reload model
model = keras.models.load_model(filepath = ...)

hasattr(model, 'normalization_mean') # False
hasattr(model, 'normalization_var')  # False

Solution 1:

this is a possibility... you can create a model subclass in this way and assign external object like not-trainable variables

X = np.random.uniform(0,1, (100,10))
y = np.random.uniform(0,1, 100)

class MyModel(tf.keras.Model):

  def __init__(self):
    super(MyModel, self).__init__()

    self.dense1 = Dense(32)
    self.dense2 = Dense(1)

  def call(self, inputs):
    x = self.dense1(inputs)
    return self.dense2(x)

model = MyModel()
model.compile('adam','mse')
model.fit(X,y)

model._normalization_mean = tf.Variable([111.], trainable=False)
model._normalization_var  = tf.Variable([222.], trainable=False)

model.save('abc.tf', save_format='tf')
model = tf.keras.models.load_model(filepath = 'abc.tf')

after loading the model you can call

model._normalization_mean.numpy()
# array([111.], dtype=float32)

here the running notebook

to save and load subclass model you can refer to this

Solution 2:

I just came across Keras preprocessing layers whose purpose seem to be exactly what you're describing.

The Keras preprocessing layers API allows developers to build Keras-native input processing pipelines. These input processing pipelines can be used as independent preprocessing code in non-Keras workflows, combined directly with Keras models, and exported as part of a Keras SavedModel.

With Keras preprocessing layers, you can build and export models that are truly end-to-end: models that accept raw images or raw structured data as input; models that handle feature normalization or feature value indexing on their own.