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>]