[FIXED] Saving Masking layer with functional API to .keras file

Issue

I'm using a custom mask in my Keras model. When I try to load the model with model=tf.saved_model.load('model.keras') from a .keras file, I get the following error:

TypeError: <keras.src.layers.core.masking.Masking object at 0x7e7735fa1460> 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.

Here is my model:

N_FEATURES = 6                                                                                       
MASK_VALUE = np.asarray([0.0 for i in range(N_FEATURES)])                                            
def get_clean_model():                                                                   
    # Input layer                                                                                    
    input_layer = Input(shape=(None, N_FEATURES))                                                            
    masked_input = Masking(mask_value=MASK_VALUE)(input_layer)                                       
                                                                                                 
    # LSTM layer with regularization                                                                 
    lstm_layer = LSTM(units=N_FEATURES, activation='tanh', return_sequences=True,                    
                  recurrent_regularizer='l2', kernel_regularizer='l2')(masked_input)             
                                                                                                 
    # Dropout layer                                                                                  
    dropout_layer = Dropout(0.05)(lstm_layer)                                                        

    # Dense layers with regularization                                                               
    dense_layer1 = Dense(N_FEATURES, activation='sigmoid', kernel_regularizer='l2')(dropout_layer)   
                                                                                                 
    # Skip connection: Concatenate masked input with dense_layer1                                    
    concatenated_layer = Concatenate()([masked_input, dense_layer1])                                 
    dense_layer2 = Dense(N_FEATURES*2, activation='sigmoid', kernel_regularizer='l2')(dense_layer1)  
                                                                                                 
    # Dropout layer                                                                                  
    dropout_layer2 = Dropout(0.05)(dense_layer2)                                                     
                                                                                                 
    # Output layer                                                                                   
    output_layer = Dense(1, activation='sigmoid')(dropout_layer2)                                    
                                                                                                 
    # Create the model                                                                               
    model = Model(inputs=input_layer, outputs=output_layer)

This question shows how to do it for custom layers, but I don't have any custom layers. How can I serialize and retrieve this model? Thank you!

Solution

The problem might be caused by new version of wrapt package, for more context see here

Could you try with setting this environment variable?

WRAPT_DISABLE_EXTENSIONS=true

Answered By - Gagik