Respect Keras layer names for output operations in Concrete Function

[DLumi]

Click to expand!

Issue Type

Feature Request

Have you reproduced the bug with TF nightly?

Yes

Source

binary

Tensorflow Version

2.12.0

Custom Code

Yes

OS Platform and Distribution

No response

Mobile device

No response

Python version

3.8.10

Bazel version

No response

GCC/Compiler version

No response

CUDA/cuDNN version

No response

GPU model and memory

No response

Current Behaviour?

When converting a Keras model to concrete function, you can preserve the input name by creating a named TensorSpec, but the outputs are always created for you by just slapping tf.identity on top of whatever you had there, even if it was a custom named tf.identity operation. Since many converters rely on concrete functions to make their own representation (TFLite, ONNX, CoreML, etc), this behavior messes up the output operation names, often making them inconsistent with each other.
There's currently no workaround for that. You can access previous graph nodes by calling a layer named like {model_name}/{output_layer_name} when doing inference on frozen graph itself, but it won't help you in any way to convert the model.

So I'd be happy to see one of those things as a solution to that:

1. Add an option to explicitly specify the TensorSpec for outputs, just the way we do it for inputs. This would be the most obvious and convenient way of doing it from a user standpoint
2. Don't add new identity operations on top of existing ones. More of a kludge, but would get the job done
3. Add an option to rename operations in concrete function post factum.
4. Add an option to cut off the operations in concrete function past a certain node.
5. Add an option to convert a graph into a concrete function. Since you can directly modify graphs, this could work as well

Standalone code to reproduce the issue

import tensorflow as tf

# Build simple model
inputs = tf.keras.Input((224, 224, 3), name='custom_input_layer')
x = tf.keras.layers.Flatten()(inputs)
x = tf.keras.layers.Dense(512, activation='relu')(x)
x = tf.keras.layers.Dense(256, activation='relu')(x)
x = tf.keras.layers.Dense(128, activation='relu')(x)
x = tf.keras.layers.Dense(1, activation='sigmoid', name='custom_output_layer')(x)

model = tf.keras.Model(inputs=inputs, outputs=x, name='my_custom_model_name')
model.summary()

input_tensors = [tf.TensorSpec(shape=inp.shape, dtype=tf.float32, name=inp.name) for inp in model.inputs]
concrete_function = tf.function(lambda x: model(x)).get_concrete_function(x=input_tensors)

print(concrete_function.inputs)  # we can see 'custom_input_layer:0' is there. So is the 'true' output 'my_custom_model_name/custom_output_layer/BiasAdd/ReadVariableOp/resource:0'
print(concrete_function.outputs)  # pesky Identity node gets inserted

Relevant log output

Model: "my_custom_model_name"
_________________________________________________________________
 Layer (type)                Output Shape              Param #   
=================================================================
 custom_input_layer (InputLa  [(None, 224, 224, 3)]    0         
 yer)                                                            
                                                                 
 flatten (Flatten)           (None, 150528)            0         
                                                                 
 dense (Dense)               (None, 512)               77070848  
                                                                 
 dense_1 (Dense)             (None, 256)               131328    
                                                                 
 dense_2 (Dense)             (None, 128)               32896     
                                                                 
 custom_output_layer (Dense)  (None, 1)                129       
                                                                 
=================================================================
Total params: 77,235,201
Trainable params: 77,235,201
Non-trainable params: 0
_________________________________________________________________
[<tf.Tensor 'custom_input_layer:0' shape=(None, 224, 224, 3) dtype=float32>, <tf.Tensor 'my_custom_model_name/dense/MatMul/ReadVariableOp/resource:0' shape=() dtype=resource>, <tf.Tensor 'my_custom_model_name/dense/BiasAdd/ReadVariableOp/resource:0' shape=() dtype=resource>, <tf.Tensor 'my_custom_model_name/dense_1/MatMul/ReadVariableOp/resource:0' shape=() dtype=resource>, <tf.Tensor 'my_custom_model_name/dense_1/BiasAdd/ReadVariableOp/resource:0' shape=() dtype=resource>, <tf.Tensor 'my_custom_model_name/dense_2/MatMul/ReadVariableOp/resource:0' shape=() dtype=resource>, <tf.Tensor 'my_custom_model_name/dense_2/BiasAdd/ReadVariableOp/resource:0' shape=() dtype=resource>, <tf.Tensor 'my_custom_model_name/custom_output_layer/MatMul/ReadVariableOp/resource:0' shape=() dtype=resource>, <tf.Tensor 'my_custom_model_name/custom_output_layer/BiasAdd/ReadVariableOp/resource:0' shape=() dtype=resource>]
[<tf.Tensor 'Identity:0' shape=(None, 1) dtype=float32>]

[sushreebarsa]

@DLumi
This issue seems to be Keras issue. Please post this issue on keras-team/keras repo. as Keras development is fully moving to github.com/keras-team/keras. All issues and PRs related to keras will be addressed in that repo.
To know more see this TF forum discussion ;
https://discuss.tensorflow.org/t/keras-project-moved-to-new-repository-in-https-github-com-keras-team-keras/1999
Thank you!

[DLumi]

I really don't see how it is the Keras issue, since it seems like a manifestation of default behavior for calling tf.function().get_concrete_function(). If you don't specify tensor names in provided list of TensorSpecs, the algorithm would just stick to defaults (x for inputs; Identity for outputs), but if you do name your tensors, the x would be replaced with whatever you provided. Sadly that's not the case for outputs that get named Identity, since the algorithm slaps a fresh copy of default Identity operation on top of whatever there was before, and you simply don't have an option to provide a spec for this, as outputs seem to be tracked and named automatically. Thus you have no ability of preserving your Keras output names.
I will repost this in Keras repo as you suggested, but I'll probably get redirected back here pretty quickly.

[SuryanarayanaY]

Hi @DLumi,

As the issue is already on Keras repo can we close it here so that it can be better tracked at single repo. Thanks!

[DLumi]

Hi @DLumi,

As the issue is already on Keras repo can we close it here so that it can be better tracked at single repo. Thanks!

Honestly, I'd keep both issues open in case there are some tweaks required from the TF team. Unless the Keras team explicitly says they will handle this themselves, that is.
But I understand your wish to reduce issues piling up, so feel free to close this if that's really needed.

[DLumi]

@SuryanarayanaY as expected, I got bounced back to you guys.
keras-team/tf-keras#219