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

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?

CodePudding user response：

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.