Saving and loading custom Keras models with custom RNN Layer

General Discussion
, , ,
#1

I subclassed a Keras layer from tf.keras.layers.AbstractRNNCell and created a custom RNN layer with a set of weights and states. I then passed this into a custom Keras model class (subclassed from tf.keras.Models). However, after training the model, when I try to go model.save('my_model', save_format='tf' ) and `tf.keras.models.load_model(‘my_model’)', I get the following error

ValueError: Unable to restore custom object of type _tf_keras_rnn_layer. Please make sure that any custom layers are included in the custom_objects arg when calling load_model() and make sure that all layers implement get_config and from_config.

#3

@Lu_Bin_Liu,

Load the model and include the custom RNN layer in the custom_objects argument

In the custom_objects argument, you need to provide a dictionary with the name of the custom layer as the key and the actual layer class as the value.

Thank you!

#5

Hi, would you mind to provide a snippet of your code when you do implement this solution? Thanks!

#6

@R_M,

Welcome to the Tensorflow Forum!

Sure, sample code is shown below

# Load the model from the saved file
loaded_model = tf.keras.models.load_model("custom_model", custom_objects={"CustomRNN": CustomRNN})

To know more please refer to sample example here how_savedmodel_handles_custom_objects.

Thank you!