如何为Keras Functional API模型的输出指定自定义名称？

使用Keras Functional API开发多输出模型后，用saved_model_cli查看SavedModel的签名定义，发现输出属性被默认命名为output_0、output_1、output_2，但预期是t1、t2、t3（与代码中输出层的名称一致）。

saved_model_cli输出结果 #

$ 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

模型代码 #

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)

希望继续使用Functional API，而非继承tf.Model类，如何指定输出属性的名称？

有两种可靠的方法可以在Functional API中保留输出名称：

方法1：自定义导出签名 #

在保存模型时，用tf.function包装推理逻辑并显式指定输入输出签名，确保输出字典的键被正确映射到SavedModel的输出名称：

import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers, models

# （模型定义代码不变）

# 定义带输入签名的推理函数
@tf.function(input_signature=[
    {
        'f1': tf.TensorSpec(shape=(None, 1), dtype=tf.float32),
        'f2': tf.TensorSpec(shape=(None, 1), dtype=tf.float32),
        'f3': tf.TensorSpec(shape=(None, 1), dtype=tf.float32),
        'f4': tf.TensorSpec(shape=(None, 1), dtype=tf.float32)
    }
])
def serving_fn(inputs):
    return model(inputs)

# 保存模型时指定自定义签名
tf.saved_model.save(
    model,
    export_dir="/serving_model_folder/new_version",
    signatures={'serving_default': serving_fn}
)

方法2：为输出张量显式命名 #

通过tf.identity为每个输出张量赋予明确名称，确保SavedModel导出时保留该名称：

# 修改输出层定义部分
output_layers = {
    't1': tf.identity(layers.Dense(1, activation='sigmoid', name='t1')(x), name='t1'),
    't2': tf.identity(layers.Dense(1, activation='sigmoid', name='t2')(x), name='t2'),
    't3': tf.identity(layers.Dense(1, activation='sigmoid', name='t3')(x), name='t3'),
}

model = models.Model(input_layers, output_layers)

# 正常保存模型即可
model.save("/serving_model_folder/new_version")

两种方法都能让saved_model_cli输出中显示t1、t2、t3作为输出属性名，无需切换到继承tf.Model的方式。

内容的提问来源于stack exchange，提问作者Mehran