如何命名 Keras Functional API 模型的输出？

Question

我有一个使用 Keras 开发的 ML 模型，更准确地说，它使用的是函数式 API。保存模型并使用 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

如您所见，3 个输出属性被命名为：output_0、output_1 和output_2。这就是我实例化模型的方式：

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)

我希望保存的模型将输出属性命名为 t1、t2 和 t3。 Searching online，我发现如果我将我的模型子类化为 tf.Model 类，我可以重命名它们：

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})

但我想继续使用 Functional API。在使用 Keras Functional API 时，有什么方法可以指定输出属性的名称吗？

Answer 1

我设法以不同的方式解决了这个问题。它依赖于签名并添加一个新层只是为了重命名张量。

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)

此代码的要点是 CustomModuleWithOutputName 类。它是 Keras 的 Layer 的子类，它所做的只是为输出索引命名。该层在保存之前添加到 serving_default 签名中的模型图中。这是一个有点愚蠢的解决方案，但它确实有效。此外，它依赖于原始函数式 API 返回的张量的顺序。

我希望我原来的方法能奏效。但既然没有，至少我有这个来买单。

Answer 2

参考链接： https://github.com/tensorflow/tensorflow/issues/48157

我们可以定义模型：

tf.keras.Model({"inputs": input_tensor}, {"boxes": boxes, "scores": scores})

并在保存的模型签名中设置输出名称：

@tf.function(input_signature=[tf.TensorSpec([None, 32, 32, 128], dtype=tf.float32, name="input")])
def override_output_signatures(input_tensor):
   outputs = model(input_tensor)
   return {"boxes": outputs["boxes"], "scores": outputs["scores"]}

“boxes”和“scores”是输出名称。