I am trying to save a Keras model in a H5 file. The Keras model has a custom layer . When I try to restore the model , I get the following error:
---------------------------------------------------------------------------
ValueError Traceback (most recent call last)
<ipython-input-5-0fbff9b56a9d> in <module>()
1 model.save('model.h5')
2 del model
----> 3 model = tf.keras.models.load_model('model.h5')
8 frames
/usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/utils/generic_utils.py in class_and_config_for_serialized_keras_object(config, module_objects, custom_objects, printable_module_name)
319 cls = get_registered_object(class_name, custom_objects, module_objects)
320 if cls is None:
--> 321 raise ValueError('Unknown ' + printable_module_name + ': ' + class_name)
322
323 cls_config = config['config']
ValueError: Unknown layer: CustomLayer
Could you please tell me how I am supposed to save and load weights of all the custom Keras layers too? (Also, there was no warning when saving, will it be possible to load models from H5 files which I have already saved but can't load back now?)
Here is the minimal working code sample (MCVE) for this error, as well as the full expanded message: Google Colab Notebook
Just for completeness, this is the code I used to make my custom layer. get_config
and from_config
are both working fine.
class CustomLayer(tf.keras.layers.Layer):
def __init__(self, k, name=None):
super(CustomLayer, self).__init__(name=name)
self.k = k
def get_config(self):
return {'k': self.k}
def call(self, input):
return tf.multiply(input, 2)
model = tf.keras.models.Sequential([
tf.keras.Input(name='input_layer', shape=(10,)),
CustomLayer(10, name='custom_layer'),
tf.keras.layers.Dense(1, activation='sigmoid', name='output_layer')
])
model.save('model.h5')
model = tf.keras.models.load_model('model.h5')
Correction number 1 is to use Custom_Objects
while loading
the Saved Model
ie, replace the code,
new_model = tf.keras.models.load_model('model.h5')
with
new_model = tf.keras.models.load_model('model.h5', custom_objects={'CustomLayer': CustomLayer})
Since we are using Custom Layers
to build
the Model
and before Saving
it, we should use Custom Objects
while Loading
it.
Correction number 2 is to add **kwargs
in the __init__
function of the Custom Layer like
def __init__(self, k, name=None, **kwargs):
super(CustomLayer, self).__init__(name=name)
self.k = k
super(CustomLayer, self).__init__(**kwargs)
Complete working code is shown below:
import tensorflow as tf
class CustomLayer(tf.keras.layers.Layer):
def __init__(self, k, name=None, **kwargs):
super(CustomLayer, self).__init__(name=name)
self.k = k
super(CustomLayer, self).__init__(**kwargs)
def get_config(self):
config = super(CustomLayer, self).get_config()
config.update({"k": self.k})
return config
def call(self, input):
return tf.multiply(input, 2)
model = tf.keras.models.Sequential([
tf.keras.Input(name='input_layer', shape=(10,)),
CustomLayer(10, name='custom_layer'),
tf.keras.layers.Dense(1, activation='sigmoid', name='output_layer')
])
tf.keras.models.save_model(model, 'model.h5')
new_model = tf.keras.models.load_model('model.h5', custom_objects={'CustomLayer': CustomLayer})
print(new_model.summary())
Output of the above code is shown below:
WARNING:tensorflow:No training configuration found in the save file, so the model was *not* compiled. Compile it manually.
Model: "sequential_1"
_________________________________________________________________
Layer (type) Output Shape Param #
=================================================================
custom_layer_1 (CustomLayer) (None, 10) 0
_________________________________________________________________
output_layer (Dense) (None, 1) 11
=================================================================
Total params: 11
Trainable params: 11
Non-trainable params: 0
Hope this helps. Happy Learning!
You can provide manually the mapping custom_objects
in the load_model
method as mentioned in the answer https://stackoverflow.com/a/62326857/8056572 but it can be tedious when you have a lot of custom layers (or any custom callables defined. eg metrics, losses, optimizers, ...).
Tensorflow provides a utils function to do it automatically: tf.keras.utils.register_keras_serializable
You have to update your CustomLayer
as follows:
import tensorflow as tf
@tf.keras.utils.register_keras_serializable()
class CustomLayer(tf.keras.layers.Layer):
def __init__(self, k, **kwargs):
self.k = k
super(CustomLayer, self).__init__(**kwargs)
def get_config(self):
config = super().get_config()
config["k"] = self.k
return config
def call(self, input):
return tf.multiply(input, 2)
Here is the complete working code:
import tensorflow as tf
@tf.keras.utils.register_keras_serializable()
class CustomLayer(tf.keras.layers.Layer):
def __init__(self, k, **kwargs):
self.k = k
super(CustomLayer, self).__init__(**kwargs)
def get_config(self):
config = super().get_config()
config["k"] = self.k
return config
def call(self, input):
return tf.multiply(input, 2)
def main():
model = tf.keras.models.Sequential(
[
tf.keras.Input(name='input_layer', shape=(10,)),
CustomLayer(10, name='custom_layer'),
tf.keras.layers.Dense(1, activation='sigmoid', name='output_layer')
]
)
print("SUMMARY OF THE MODEL CREATED")
print("-" * 60)
print(model.summary())
model.save('model.h5')
del model
print()
print()
model = tf.keras.models.load_model('model.h5')
print("SUMMARY OF THE MODEL LOADED")
print("-" * 60)
print(model.summary())
if __name__ == "__main__":
main()
And the corresponding output:
SUMMARY OF THE MODEL CREATED
------------------------------------------------------------
Model: "sequential"
_________________________________________________________________
Layer (type) Output Shape Param #
=================================================================
custom_layer (CustomLayer) (None, 10) 0
_________________________________________________________________
output_layer (Dense) (None, 1) 11
=================================================================
Total params: 11
Trainable params: 11
Non-trainable params: 0
_________________________________________________________________
None
WARNING:tensorflow:No training configuration found in the save file, so the model was *not* compiled. Compile it manually.
SUMMARY OF THE MODEL LOADED
------------------------------------------------------------
Model: "sequential"
_________________________________________________________________
Layer (type) Output Shape Param #
=================================================================
custom_layer (CustomLayer) (None, 10) 0
_________________________________________________________________
output_layer (Dense) (None, 1) 11
=================================================================
Total params: 11
Trainable params: 11
Non-trainable params: 0
_________________________________________________________________
None
The other answer provided here does not work in a real example. The code example is running in the same process, so the CustomLayer
is available in-memory when loading the model. If you try to load the model in a new process it will fail and ask you to use the custom_objects
argument. I am still looking for a clean way to export custom objects from the source code.
The technical post webpages of this site follow the CC BY-SA 4.0 protocol. If you need to reprint, please indicate the site URL or the original address.Any question please contact:yoyou2525@163.com.