Keras SavedModel format doesn't support masking in Keras 3

Loading a SavedModel in Keras 3 loses any mask information on the model output. Loading a SavedModel in tf-keras preserves the mask information, as expected. Here is an example:

USE_KERAS_3 = True

import tensorflow as tf
if USE_KERAS_3:
    import keras
else:
    import tf_keras as keras



inp = keras.Input((1,))
x = keras.layers.Masking(mask_value=0.0)(inp)
model = keras.Model(inp, x)

if USE_KERAS_3:
    model.export("saved_model")
    loaded = keras.layers.TFSMLayer("saved_model")
else:
    model.save("saved_model", save_format="tf")
    loaded = keras.models.load_model("saved_model")


print(loaded(tf.zeros((1, 1)))._keras_mask)

With USE_KERAS_3=True this gives:

AttributeError: 'tensorflow.python.framework.ops.EagerTensor' object has no attribute '_keras_mask'

With USE_KERAS_3=False this gives:

tf.Tensor([False], shape=(1,), dtype=bool)

Comment From: sonali-kumari1

Hi @drasmuss -

I tried to reproduce the issue and I got AttributeError: 'tensorflow.python.framework.ops.EagerTensor' object has no attribute '_keras_mask' with latest version of keras(3.9.2). It seems using model.export to save and loading model with loaded = keras.layers.TFSMLayer("saved_model") results in the mask information being lost.

However, I tried to save the model using model.save("saved_model.keras") and loaded the model using loaded = keras.models.load_model("saved_model.keras") and I was able to get the following output tensor, without any error:

tf.Tensor([False], shape=(1,), dtype=bool)

Attaching gist for your reference. Thanks!

Comment From: drasmuss

Unfortunately the model I am working with uses custom layers, and I need to be able to load the model in an environment where the original source code isn't available, so I have to use the SavedModel format, not the Keras format.

Comment From: sonali-kumari1

@drasmuss, If you are loading model with loaded = keras.layers.TFSMLayer("saved_model"), please note that the reloaded object retains none of the internal structure or custom methods of the original object. Since you are using custom layers, I would suggest implementing masking logic directly in your call() method to ensure it becomes a part of the saved computation graph. Also, make sure to implement get_config() and from_config() methods for your custom layers. Please refer to this documentation for more details on TFSMLayer. Thanks!

Comment From: drasmuss

The SavedModel saving and loading works in tf-keras, so I don't think it's a problem with the custom layer code. Also note that the example above doesn't use any custom layers. The issue is specifically with Keras 3, and how it implements the saving/loading.

Comment From: divyashreepathihalli

@drasmuss, having your build() method and get_config() and from_config(), should be accurately implemented for saving and loading to be accurate. Yes TF backend is more forgiving for saving and loading and other backends are not.

Comment From: drasmuss

As mentioned, in the example above I'm using a standard, built-in Keras layer (keras.layers.Masking), so it has nothing to do with my implementation.

Comment From: hertschuh

@drasmuss ,

Unfortunately the model I am working with uses custom layers, and I need to be able to load the model in an environment where the original source code isn't available, so I have to use the SavedModel format, not the Keras format.

Can you elaborate on your use case? Is this for serving/inference?

You say the original source code isn't available. Is the Keras source code available?

This:

loaded = keras.models.load_model("saved_model")

Requires the source code to be available. It does some magic to reconnect the saved model to the Python code.

However, if you do:

loaded = tf.saved_model.load("saved_model_keras2").signatures["serving_default"]

This does not require the source code, but the masks won't work either.

The pure graph representation of the exported function in a "saved model" does not support masks. Doing model.save("saved_model", save_format="tf") adds some information in addition to the "saved model" graph, which is then used by keras.models.load_model("saved_model") to reconnect the graph to the code. But that requires the code.

The workaround is to return masks explicitly. With this, your saved model will return a dict with "output" and "mask".

@tf.function
def fn(x):
    output = model(x)
    return {"output": output, "mask": output._keras_mask}

export_archive = keras.export.ExportArchive()
export_archive.track(model)
export_archive.add_endpoint(
    "serve",
    fn,
    input_signature=[tf.TensorSpec(shape=(1, 1), dtype=tf.float32)],
)
export_archive.write_out("saved_model")

loaded = keras.layers.TFSMLayer("saved_model")
print(loaded(tf.zeros((1, 1))))

{'mask': <tf.Tensor: shape=(1,), dtype=bool, numpy=array([False])>, 'output': <tf.Tensor: shape=(1, 1), dtype=float32, numpy=array([[0.]], dtype=float32)>}

If you want to get fancy, you can then create a wrapper to recombine them:

def loaded_wrapper(input):
    output_dict = loaded(input)
    output = output_dict["output"]
    output._keras_mask = output_dict["mask"]
    return output

print(loaded_wrapper(tf.zeros((1, 1)))._keras_mask)

tf.Tensor([False], shape=(1,), dtype=bool)

Comment From: github-actions[bot]

This issue is stale because it has been open for 14 days with no activity. It will be closed if no further activity occurs. Thank you.

Comment From: github-actions[bot]

This issue was closed because it has been inactive for 28 days. Please reopen if you'd like to work on this further.