model.export() fails to save output layer name in the SavedModel about keras

[mehtamansi29]

Hi @chriscarollo -

Thanks for reporting the issue. I have tested the code snippet and reproduces the reported behaviour. Attached gist file for reference.

We will look into the issue and update you the same.

[chriscarollo]

Trivially reproduced:

input = keras.layers.Input( (1,), name='input_1' )
output = keras.layers.Dense( 8, name='output_1' )( input )
m = keras.Model( inputs=[input], outputs=[output], name='model_1' )
m.compile()
m.summary()
Model: "model_1"
┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━┓
┃ Layer (type)                         ┃ Output Shape                ┃         Param # ┃
┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━┩
│ input_1 (InputLayer)                 │ (None, 1)                   │               0 │
├──────────────────────────────────────┼─────────────────────────────┼─────────────────┤
│ output_1 (Dense)                     │ (None, 8)                   │              16 │
└──────────────────────────────────────┴─────────────────────────────┴─────────────────┘
 Total params: 16 (64.00 B)
 Trainable params: 16 (64.00 B)
 Non-trainable params: 0 (0.00 B)
>>> m.export( 'test' )
INFO:tensorflow:Assets written to: test/assets
INFO:tensorflow:Assets written to: test/assets
Saved artifact at 'test'. The following endpoints are available:

* Endpoint 'serve'
  args_0 (POSITIONAL_ONLY): TensorSpec(shape=(None, 1), dtype=tf.float32, name='input_1')
Output Type:
  TensorSpec(shape=(None, 8), dtype=tf.float32, name=None)
Captures:
  140410840896736: TensorSpec(shape=(), dtype=tf.resource, name=None)
  140410840896912: TensorSpec(shape=(), dtype=tf.resource, name=None)

[chriscarollo]

FWIW I'm running Tensorflow 2.16.1 with Keras 3.3.3, but it does seem like this repros with Tensorflow 2.15 as well.

[grasskin]

Would you mind taking a look @hertschuh? CC: @nkovela1

[hertschuh]

Hi @chriscarollo ,

Thanks for the report. I'm not really clear about what Triton needs from the saved model, but the outputs appear to have a name (more on that below).

but when I model.export() it says:

Output Type:
TensorSpec(shape=(None, 1), dtype=tf.float32, name=None)

Apparently, the only thing that is relevant in this message is the "output type". Unfortunately (and I don't know why) the name is missing.

If I also save as a .keras file, it correctly saves that layer with my "output_1" name.

Correct. Layers and outputs are very different concepts though. Do you need the layer? Or do you care about the output?

It appears that the output is named. It's just numbered starting from zero to support multiple outputs. So in your case, the output is named output_0 (independent of the layer name).

Here's how you can find out by adding this to your code above:

# Reload the model
loaded = tf.saved_model.load('test')
print("Outputs", loaded.signatures['serve'].structured_outputs.items())

Which prints

Outputs dict_items([('output_0', TensorSpec(shape=(None, 8), dtype=tf.float32, name='output_0'))])

But overall, I'm surprised Triton needs more than the name of the function, which is serve by default.