Warning : Custom mask layers require a config and must override when saving the model in keras

Question

I created a model from sequential. when I saved it I got this warning message

home/anaconda3/lib/python3.8/site-packages/tensorflow/python/keras/utils/generic_utils.py:494: CustomMaskWarning: Custom mask layers require a config and must override get_config. When loading, the custom mask layer must be passed to the custom_objects argument.
  warnings.warn('Custom mask layers require a config and must override 

I tested one image and the prediction was good the I saved my model when I loaded it again it started giving me wrong values and the prediction was all wrong.what is the correct way to say the model and load it

import numpy as np 
import matplotlib.pyplot as plt
import glob
import cv2
import os
from tensorflow import keras

from tensorflow.keras.layers import Conv2D, MaxPooling2D
from tensorflow.keras.layers import Input, Dropout, Flatten, Dense

from tensorflow.keras.layers import UpSampling2D
from tensorflow.keras.models import Model
from tensorflow.keras.layers import BatchNormalization
from tensorflow.keras.models import Sequential

input_shape = (3,1134,1134,3)
base_model = tf.keras.applications.ResNet50(
include_top=False,
weights="imagenet",
input_shape=(1134,1134,3),
pooling=max,  
)
for layer in base_model.layers[:-4]:
    layer.trainable = False
model = Sequential()
model.add(Conv2D(3,(3,3),activation='relu',padding='same'))
model.add(base_model)
model.add(Conv2D(3,(3,3),activation='relu',padding='same'))
# model.add(Convolution2D(3,(4,4),activation='relu',padding='same'))
model.add(UpSampling2D(size =(16,16)))
model.add(UpSampling2D())
model.add(BatchNormalization())
model.add(Conv2D(3,(3,3),activation='relu',padding='same'))
model.build(input_shape)
model.summary()

this is how I save it

model.save("/media/TOSHIBA EXT/trained_model/UAV_01.h5")

enter code here

model=keras.models.load_model(
    "/media/TOSHIBA EXT/trained_model/UAV_01.h5")

Answer 1

@user123 Agree with you that it was an issue with old versions (from TF2.5, TF2.6, and TF2.7).

This was resolved in recent tf-nightly. Here is a gist for reference. If you want to use a stable version, then it will be available in the upcoming TF2.8 in near future. Thanks!

Answer 2

Two other approaches to try:

1. Save the model using a directory path instead of with a path that ends in the .h5 extension. Under the hood, save does different things if you send a path that ends with .h5. If you send a directory, it will use the newer SavedModel format. You can then directly load the model with:

from tensorflow.keras.models import load_model


new_model = load_model('<path to directory used in save>')

https://www.tensorflow.org/guide/saved_model#the_savedmodel_format_on_disk

2. Only save the weights. To load the model, use your current model creation code to instantiate the object, then load the weights into the empty model object. This used to be the only way!

model.save_weights('my_model_weights.h5')
...
new_model = <build your model with your model building code>
new_model.load_weights('my_model_weights.h5')

Ref: https://keras.io/getting_started/faq/. Weights only saving

In the FAQ, there are also several other suggestions for how to handle it, but for your situation these two will probably cover it.

---

Here is a good snippet to add after your train code to make sure that the export works and is not the cause of a performance issue at inference:

# From: https://www.tensorflow.org/guide/keras/save_and_serialize#whole-model_saving_loading
# Train the model.
test_input = np.random.random((128, 32))
test_target = np.random.random((128, 1))
model.fit(test_input, test_target)

# Calling `save('my_model')` creates a SavedModel folder `my_model`.
model.save("my_model")

# It can be used to reconstruct the model identically.
reconstructed_model = keras.models.load_model("my_model")

# Let's check:
np.testing.assert_allclose(
    model.predict(test_input), reconstructed_model.predict(test_input)
)

# The reconstructed model is already compiled and has retained the optimizer
# state, so training can resume:
reconstructed_model.fit(test_input, test_target)