[Python] Why the keras.models.save_model() method does not save the weigths when the model is...

I am trying to train, save and load back a model defined in a custom class inherited from keras.Model. The definition and training is working fine. I save the trained model in the '.keras' format. When I load back the model, the architecture is preserved but the weights of each layer are missing. Here is a code example reproducing the issue :

import tensorflow as tf
import numpy as np

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

def __init__(self, dense_size, **kwargs):

self.dense_size = dense_size
super().__init__(**kwargs)


def build(self):
self.layer1 = tf.keras.layers.Dense(self.dense_size)
self.layer2 = tf.keras.layers.Dense(self.dense_size)

def call(self, inputs):
x = self.layer1(inputs)
out = self.layer2(x)
return out


def get_config(self):
config = super().get_config()
config.update({
'dense_size': self.dense_size
})
return config


input = tf.random.normal((100,100,1))
output = tf.random.normal((100,100,1))

model = dummyModel(10)

model.compile(optimizer='adam', loss='mse')
model.fit(input, output, epochs=10, batch_size=10)

tf.keras.models.save_model(model, 'dummy.keras')

modelSaved = tf.keras.models.load_model('dummy.keras')

assert model.to_json() == modelSaved.to_json(), "Model architectures are different"

for layer1, layer2 in zip(model.layers, modelSaved.layers):
weights1 = layer1.get_weights()
weights2 = layer2.get_weights()
if weights1 != []:
assert weights2, f"{layer2.name} of loaded model has empty weight list"
for w1, w2 in zip(weights1, weights2):
assert np.array_equal(w1, w2), f"Weigths of {layer1.name} are different."

print("Models are identical (architecture and weigths).")


The code throws this assertion error :

AssertionError Traceback (most recent call last)

<ipython-input-10-3ff24dc35a83> in <cell line: 46>()
48 weights2 = layer2.get_weights()
49 if weights1 != []:
---> 50 assert weights2, f"{layer2.name} of loaded model has empty weight list"
51 for w1, w2 in zip(weights1, weights2):
52 assert np.array_equal(w1, w2), f"Weigths of {layer1.name} are different."

AssertionError: dense_2 of loaded model has empty weight list


I tried to use a keras.Sequential model or the Functional API and the weights are saved properly. When the model is defined in a custom class with the appropriate get_config() method, no weights are saved. I am relatively new the keras framework, I might miss a key concept here.

Here is a collab link with the same code :

https://colab.research.google.com/drive/1H47N77k5skn9CXp4ijx3sMhGy2EL5iAd#scrollTo=0xy_DDNNO2So

Continue reading...