How to name the outputs of a Keras Functional API model? - tensorflow

Question

I have an ML model developed using Keras and more accurately, it's using Functional API. Once I save the model and use the saved_model_cli tool on it:
$ saved_model_cli show --dir /serving_model_folder/1673549934 --tag_set serve --signature_def serving_default
2023-01-12 10:59:50.836255: I tensorflow/core/util/util.cc:169] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`.
The given SavedModel SignatureDef contains the following input(s):
inputs['f1'] tensor_info:
dtype: DT_FLOAT
shape: (-1, 1)
name: serving_default_f1:0
inputs['f2'] tensor_info:
dtype: DT_FLOAT
shape: (-1, 1)
name: serving_default_f2:0
inputs['f3'] tensor_info:
dtype: DT_FLOAT
shape: (-1, 1)
name: serving_default_f3:0
inputs['f4'] tensor_info:
dtype: DT_FLOAT
shape: (-1, 1)
name: serving_default_f4:0
The given SavedModel SignatureDef contains the following output(s):
outputs['output_0'] tensor_info:
dtype: DT_FLOAT
shape: (-1)
name: StatefulPartitionedCall_1:0
outputs['output_1'] tensor_info:
dtype: DT_FLOAT
shape: (-1)
name: StatefulPartitionedCall_1:1
outputs['output_2'] tensor_info:
dtype: DT_FLOAT
shape: (-1)
name: StatefulPartitionedCall_1:2
Method name is: tensorflow/serving/predict
As you can see, the 3 output attributes are named: output_0, output_1, and output_2. This is how I'm instantiating my model:
input_layers = {
'f1': Input(shape=(1,), name='f1'),
'f2': Input(shape=(1,), name='f2'),
'f3': Input(shape=(1,), name='f3'),
'f4': Input(shape=(1,), name='f4'),
}
x = layers.concatenate(input_layers.values())
x = layers.Dense(32, activation='relu', name="dense")(x)
output_layers = {
't1': layers.Dense(1, activation='sigmoid', name='t1')(x),
't2': layers.Dense(1, activation='sigmoid', name='t2')(x),
't3': layers.Dense(1, activation='sigmoid', name='t3')(x),
}
model = models.Model(input_layers, output_layers)
I was hoping that the saved model would name the output attributes t1, t2, and t3. Searching online, I see that I can rename them if I subclass my model off tf.Model class:
class CustomModuleWithOutputName(tf.Module):
def __init__(self):
super(CustomModuleWithOutputName, self).__init__()
self.v = tf.Variable(1.)
#tf.function(input_signature=[tf.TensorSpec([], tf.float32)])
def __call__(self, x):
return {'custom_output_name': x * self.v}
module_output = CustomModuleWithOutputName()
call_output = module_output.__call__.get_concrete_function(tf.TensorSpec(None, tf.float32))
module_output_path = os.path.join(tmpdir, 'module_with_output_name')
tf.saved_model.save(module_output, module_output_path,
signatures={'serving_default': call_output})
But I would like to keep using the Functional API. Is there any way to specify the name of the output attributes while using Keras Functional API?

Answer 1

I managed to pull this off a different way. It relies on the signature and adds a new layer just to rename the tensors.
from tensorflow.keras import layers
class CustomModuleWithOutputName(layers.Layer):
def __init__(self):
super(CustomModuleWithOutputName, self).__init__()
def call(self, x):
return {'t1': tf.identity(x[0]),
't2': tf.identity(x[1]),
't3': tf.identity(x[2]),}
def _get_tf_examples_serving_signature(model):
#tf.function(input_signature=[tf.TensorSpec(shape=[None, 1], dtype=tf.float32, name='f1'),
tf.TensorSpec(shape=[None, 1], dtype=tf.float32, name='f2'),
tf.TensorSpec(shape=[None, 1], dtype=tf.float32, name='f3'),
tf.TensorSpec(shape=[None, 1], dtype=tf.float32, name='f4'),])
def serve_tf_examples_fn(f1, f2, f3, f4):
"""Returns the output to be used in the serving signature."""
inputs = {'f1': f1, 'f2': f2, 'f3': f3, 'f4': f4}
outputs = model(inputs)
return model.naming_layer(outputs)
return serve_tf_examples_fn
# This is the same model mentioned in the question (a Functional API model)
model = get_model()
# Any property name will do as long as it is not reserved
model.naming_layer = CustomModuleWithOutputName()
signatures = {
'serving_default': _get_tf_examples_serving_signature(model),
}
model.save(output_dir, save_format='tf', signatures=signatures)
The takeaway from this code is the CustomModuleWithOutputName class. It's a subclass of Keras' Layer and all it does is give names to the output indices. This layer is added to the model's graph in the serving_default signature before it is saved. It's a kinda stupid solution but it works. Also, it relies on the order of the tensors returned by the original functional API.
I was hoping my original approach would work. But since it doesn't, at least I have this one to foot the bill.