The signatures in SavedModel do not contain serving_default when a subclass of the keras model has multiple inputs

[ian-huu]

Issue type

Bug

Have you reproduced the bug with TensorFlow Nightly?

Yes

Source

binary

TensorFlow version

2.16.1

Custom code

Yes

OS platform and distribution

Linux Ubuntu 22.04.1

Mobile device

No response

Python version

3.10.13

Bazel version

No response

GCC/compiler version

No response

CUDA/cuDNN version

No response

GPU model and memory

No response

Current behavior?

The code below illustrates that when a Keras model has multiple inputs, the tf.saved_model does not function correctly, resulting in an empty signatures field.

Standalone code to reproduce the issue

import tempfile

import keras
import tensorflow as tf

path = tempfile.mkdtemp()

model1 = keras.models.Sequential([keras.layers.Dense(8), keras.layers.Dense(1)])
model1.build((None, 16))
tf.saved_model.save(model1, f"{path}/model1")
loaded1 = tf.saved_model.load(f"{path}/model1")
print(f"loaded1: {loaded1.signatures}")


class Model2(keras.Model):
    def __init__(self):
        super().__init__()

        self.fc1 = keras.layers.Dense(8)
        self.fc2 = keras.layers.Dense(1)

    def call(self, inputs, **kwargs):
        out1 = self.fc1(inputs)
        out2 = self.fc2(out1)
        return out2


model2 = Model2()
model2.build((None, 16))
tf.saved_model.save(model2, f"{path}/model2")
loaded2 = tf.saved_model.load(f"{path}/model2")
print(f"loaded2: {loaded2.signatures}")


class Model3(keras.Model):
    def __init__(self):
        super().__init__()

        self.fc1 = keras.layers.Dense(8)
        self.fc2 = keras.layers.Dense(1)

    def build(self, input_shape):
        self.fc1.build((None, sum([shape[1] for shape in input_shape])))
        self.fc2.build((None, 8))

    def call(self, inputs, **kwargs):
        concat_inputs = keras.ops.concatenate(inputs, axis=1)
        out1 = self.fc1(concat_inputs)
        out2 = self.fc2(out1)
        return out2


model3 = Model3()
model3.build([(None, 8), (None, 8)])
tf.saved_model.save(model3, f"{path}/model3")

loaded3 = tf.saved_model.load(f"{path}/model3")
print(f"loaded3: {loaded3.signatures}")

Relevant log output

loaded1: _SignatureMap({'serving_default': <ConcreteFunction (*, inputs: TensorSpec(shape=(None, 16), dtype=tf.float32, name='inputs')) -> Dict[['output_0', TensorSpec(shape=(None, 1), dtype=tf.float32, name='output_0')]] at 0x7A4E8429C8E0>})
python3.10/site-packages/keras/src/layers/layer.py:361: UserWarning: `build()` was called on layer 'model2', however the layer does not have a `build()` method implemented and it looks like it has unbuilt state. This will cause the layer to be marked as built, despite not being actually built, which may cause failures down the line. Make sure to implement a proper `build()` method.
  warnings.warn(
loaded2: _SignatureMap({'serving_default': <ConcreteFunction (*, inputs: TensorSpec(shape=(None, 16), dtype=tf.float32, name='inputs')) -> Dict[['output_0', TensorSpec(shape=(None, 1), dtype=tf.float32, name='output_0')]] at 0x7A4E841D4A00>})
loaded3: _SignatureMap({})

[sushreebarsa]

@ian-huu One workaround could be to specify the input names during saving using the tf.saved_model.save API with the input argument. This ensures that all your model's inputs are included in the serving_default signature.
Thank you!

[ian-huu]

@ian-huu One workaround could be to specify the input names during saving using the tf.saved_model.save API with the input argument. This ensures that all your model's inputs are included in the serving_default signature. Thank you!

Do you mean to do it this way? It seems to be working.

class Model3(keras.Model):
    def __init__(self):
        super().__init__()

        self.fc1 = keras.layers.Dense(8)
        self.fc2 = keras.layers.Dense(1)

    def build(self, input_shape):
        self.fc1.build((None, sum([shape[1] for shape in input_shape])))
        self.fc2.build((None, 8))

    @tf.function
    def call(self, inputs, **kwargs):
        concat_inputs = keras.ops.concatenate(inputs, axis=1)
        out1 = self.fc1(concat_inputs)
        out2 = self.fc2(out1)
        return out2


model3 = Model3()
model3.build([(None, 8), (None, 8)])
tf.saved_model.save(
    model3,
    f"{path}/model3",
    signatures=model3.call.get_concrete_function(
        (
            tf.TensorSpec(shape=(None, 8), dtype=tf.float32),
            tf.TensorSpec(shape=(None, 8), dtype=tf.float32),
        )
    ),
)

loaded3 = tf.saved_model.load(f"{path}/model3")
print(f"loaded3: {loaded3.signatures}")

# got:
# loaded3: _SignatureMap({'serving_default': <ConcreteFunction (*, inputs: TensorSpec(shape=(None, 8), dtype=tf.float32, name='inputs'), inputs_1: TensorSpec(shape=(None, 8), dtype=tf.float32, name='inputs_1')) -> Dict[['output_0', TensorSpec(shape=(None, 1), dtype=tf.float32, name='output_0')]] at 0x7F2CB1FF4190>})

However, it seems there is another problem. The loaded SavedModel does not have attributes like trainable_variables and __call__ anymore.

# for all loaded1, loaded2, loaded3
print(loaded1.trainable_variables)
# got: AttributeError: '_UserObject' object has no attribute 'trainable_variables'.

print(loaded1.__call__)
# got: AttributeError: '_UserObject' object has no attribute '__call__'.

And if I directly call model3.call, an error occurs, even though I manage to obtain the result.

print(loaded3.call((tf.zeros((1, 8)), tf.zeros((1, 8)))))

got:

tf.Tensor([[0.]], shape=(1, 1), dtype=float32)
Exception ignored in: <function AtomicFunction.__del__ at 0x72f6cf3a7760>
Traceback (most recent call last):
  File "***/site-packages/tensorflow/python/eager/polymorphic_function/atomic_function.py", line 291, in __del__
TypeError: 'NoneType' object is not subscriptable

[sushreebarsa]

@ian-huu Could you try to use tf.keras.models.load_model() to load the model and let us know? Please make sure that you are using the latest TF version?
Thank you!

[ian-huu]

@ian-huu Could you try to use tf.keras.models.load_model() to load the model and let us know? Please make sure that you are using the latest TF version? Thank you!

@sushreebarsa It seems that tf.keras.models.load_model supports only the Keras or H5 format, not the TensorFlow SavedModel format:

Traceback (most recent call last):
  File "test.py", line 13, in <module>
    loaded1 = tf.keras.models.load_model(f"{path}/model1")
  File "**/python3.10/site-packages/keras/src/saving/saving_api.py", line 193, in load_model
    raise ValueError(
ValueError: File format not supported: filepath=/tmp/**/model1. Keras 3 only supports V3 `.keras` files and legacy H5 format files (`.h5` extension). Note that the legacy SavedModel format is not supported by `load_model()` in Keras 3. In order to reload a TensorFlow SavedModel as an inference-only layer in Keras 3, use `keras.layers.TFSMLayer(/tmp/**/model1, call_endpoint='serving_default')` (note that your `call_endpoint` might have a different name).

And I tried loaded3 = keras.layers.TFSMLayer(f"{path}/model3", call_endpoint='serving_default'), but it didn't work either.

Traceback (most recent call last):
  File "test.py", line 72, in <module>
    loaded3 = keras.layers.TFSMLayer(f"{path}/model3", call_endpoint='serving_default')
  File "**/python3.10/site-packages/keras/src/export/export_lib.py", line 731, in __init__
    raise ValueError(
ValueError: The endpoint 'serving_default' is neither an attribute of the reloaded SavedModel, nor an entry in the `signatures` field of the reloaded SavedModel. Select another endpoint via the `call_endpoint` argument. Available endpoints for this SavedModel: []

[ian-huu]

I noticed that TensorFlow released version 2.17.0rc0 and gave it a try, the issue remains unresolved. 😕

[sushreebarsa]

@ian-huu There are two SavedModel formats: legacy and the new high-level format. The error message indicates you're dealing with the legacy format, which isn't compatible with load_model. Please refer this guide SavedModel format.
Thank you!