Tensorflow 2.17 - Saving custom model does not work for me

TensorFlow
1

(tensorflow 2.17, Windows 10)

I can’t save and restore a custom model using subclassing method.

Please find below the code to reproduce the problem :

import numpy as np
import tensorflow as tf

x = np.random.random((1000, 32))
y = np.random.random((1000, 1))

@tf.keras.utils.register_keras_serializable()
class MModel(tf.keras.Model):

    def __init__(self):
        super().__init__()
        self.dense1 = tf.keras.layers.Dense(10)
        self.dense2 = tf.keras.layers.Dense(1)

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

model = MModel()

model.compile( optimizer="adam",  loss="mse",  metrics=["mae"])
model.fit( x, y, epochs=5)

model.save("save1.keras")
model_ = tf.keras.models.load_model("save1.keras")

print()
print(f"{model.evaluate(x, y,verbose=0)  = }")
print(f"{model_.evaluate(x, y,verbose=0) = }")

for m, m_ in zip(model.weights, model_.weights):
    np.testing.assert_allclose(m.numpy(), m_.numpy())

Please note that the init() has no parameters, hence no need to override get_config().

The error message :

TypeError: <class '__main__.MModel'> could not be deserialized properly. Please ensure that components that are Python object instances (layers, models, etc.) returned by `get_config()` are explicitly deserialized in the model's `from_config()` method.

config={'module': None, 'class_name': 'MModel', 'config': {'trainable': True, 'dtype': {'module': 'keras', 'class_name': 'DTypePolicy', 'config': {'name': 'float32'}, 'registered_name': None}}, 'registered_name': 'Custom>MModel', 'build_config': {'input_shape': [None, 32]}, 'compile_config': {'optimizer': {'module': 'keras.optimizers', 'class_name': 'Adam', 'config': {'name': 'adam', 'learning_rate': 0.0010000000474974513, 'weight_decay': None, 'clipnorm': None, 'global_clipnorm': None, 'clipvalue': None, 'use_ema': False, 'ema_momentum': 0.99, 'ema_overwrite_frequency': None, 'loss_scale_factor': None, 'gradient_accumulation_steps': None, 'beta_1': 0.9, 'beta_2': 0.999, 'epsilon': 1e-07, 'amsgrad': False}, 'registered_name': None}, 'loss': 'mse', 'loss_weights': None, 'metrics': ['mae'], 'weighted_metrics': None, 'run_eagerly': False, 'steps_per_execution': 1, 'jit_compile': False}}.

Exception encountered: Unable to revive model from config. When overriding the `get_config()` method, make sure that the returned config contains all items used as arguments in the  constructor to <class '__main__.MModel'>, which is the default behavior. You can override this default behavior by defining a `from_config(cls, config)` class method to specify how to create an instance of MModel from its config.

Received config={'trainable': True, 'dtype': {'module': 'keras', 'class_name': 'DTypePolicy', 'config': {'name': 'float32'}, 'registered_name': None}}

Error encountered during deserialization: MModel.__init__() got an unexpected keyword argument 'trainable'
2

Hi @Gilles_Jack, When saving a model that includes custom objects, such as a subclassed Layer, you must define a get_config() method on the object class. please refer to this gist for working code example and refer to this document to know more about the saving custom model. Thank You.

3

This toy model is a custom model but does not include any custom model.

5

Hi @Gilles_Jack, Even though the model does not contain any custom code and uses default keras layers the model architecture which you are trying contains sub classed layers and the MModel is a custom Python class that defines the model architecture. so you have to define the get_config( ) method. Thank You.

6

Thanks for your help Kiran but I am still confused.

The first thing I tried is to override get_config() though it seems useless as MModel has no parameters. I got the same TypeError: <class ‘main.MModel’> could not be deserialized properly. error.

@tf.keras.utils.register_keras_serializable()
class MModel(tf.keras.Model):

def __init__(self):
    super().__init__()
    self.dense1 = tf.keras.layers.Dense(10)
    self.dense2 = tf.keras.layers.Dense(1)

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

def get_config(self):
    config = super().get_config()
    return config

However without get_config(), I get

Error encountered during deserialization: MModel.init() got an unexpected keyword argument ‘name’

and with get_config() , I get ::

Error encountered during deserialization: MModel.init() got an unexpected keyword argument ‘trainable’

I don’t think I have to deal with name or trainable parameters. And the documentation agrees with me.

Please note that Functional API case below works fine :

import numpy as np
import tensorflow as tf

x = np.random.random((1000, 32))
y = np.random.random((1000, 1))

@tf.keras.utils.register_keras_serializable()
class MModel(tf.keras.Model):
pass

inputs = tf.keras.Input(shape=(32,))
o = tf.keras.layers.Dense(10)(inputs)
o = tf.keras.layers.Dense(1)(o)
model = MModel(inputs, o)

model.compile( optimizer=“adam”, loss=“mse”, metrics=[“mae”])
model.fit( x, y, epochs=5)

model.save(“save1.keras”)
model_ = tf.keras.models.load_model(“save1.keras”)

print()
print(f"{model.evaluate(x, y,verbose=0) = }“)
print(f”{model_.evaluate(x, y,verbose=0) = }")

for m, m_ in zip(model.weights, model_.weights):
np.testing.assert_allclose(m.numpy(), m_.numpy())

8

Hi @Gilles_Jack, while saving the model that was defined using the functional API does not require the get_confing method but model that was defined using the sub classing or custom format requires get_config in order to serialize and deserialize the model while saving and loading.

while using the get_config method I did not face any error. please refer to this gist for working code. Thank You.

10

I eventually solve this issue thanks to your working code.
In init(), I forgot to pass the keyword arguments to the superclass.

The default implementation of get_config() can be removed.
So this code is working too :

import numpy as np
import tensorflow as tf

x = np.random.random((1000, 32))
y = np.random.random((1000, 1))

@tf.keras.utils.register_keras_serializable()
class MModel(tf.keras.Model):

def __init__(self, **kwargs):
    super().__init__(**kwargs)
    self.dense1 = tf.keras.layers.Dense(10)
    self.dense2 = tf.keras.layers.Dense(1)

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

model = MModel()

model.compile( optimizer=“adam”, loss=“mse”, metrics=[“mae”])
model.fit( x, y, epochs=5)

model.save(“save1.keras”)
model_ = tf.keras.models.load_model(“save1.keras”)

print()
print(f"{model.evaluate(x, y,verbose=0) = }“)
print(f”{model_.evaluate(x, y,verbose=0) = }")

for m, m_ in zip(model.weights, model_.weights):
np.testing.assert_allclose(m.numpy(), m_.numpy())

Thanks again Kiran Sai.