Tensorflow 2.0：如何在使用tf.saved_model时更改输出签名

Question

我希望更改已保存的模型的输入和输出签名，我使用tf.Module对象构建了主模型的操作。

class Generator(tf.Module):
    def __init__(....):
        super(Generator, self).__init__(name=name)
        ...       
        with self.name_scope:
             ...
    @tf.Module.with_name_scope
    def __call__(self, input):
        ...

    @tf.function
    def serve_function(self, input):
        out = self.__call__(input)
        return out



call = model.Generator.serve_function.get_concrete_function(tf.TensorSpec([None, 256, 256, 3], tf.float32))
tf.saved_model.save(model.Generator, os.path.join(train_log_dir, 'frozen'))

我正在加载模型，但是我的签名是"default_serving"和"output_0"，我该如何更改？

- Miguel Angel Alba Acosta

Answer 1

我找到一种定义输出签名的方法，而不使用tf.Module，方法是定义一个返回输出字典的tf.function，其中字典中使用的键将是输出名称。

# Create the model
model = ...

# Train the model
model.fit(...)

# Define where to save the model
export_path = "..."

@tf.function()
def my_predict(my_prediction_inputs):
   inputs = {
        'my_serving_input': my_prediction_inputs,
   }
   prediction = model(inputs)
   return {"my_prediction_outputs": prediction}

my_signatures = my_predict.get_concrete_function(
   my_prediction_inputs=tf.TensorSpec([None,None], dtype=tf.dtypes.float32, name="my_prediction_inputs")
)

# Save the model.
tf.saved_model.save(
    model,
    export_dir=export_path,
    signatures=my_signatures
)

这会生成以下签名：

signature_def['serving_default']:
  The given SavedModel SignatureDef contains the following input(s):
    inputs['my_prediction_inputs'] tensor_info:
        dtype: DT_FLOAT
        shape: (-1, -1)
        name: serving_default_my_prediction_inputs:0
  The given SavedModel SignatureDef contains the following output(s):
    outputs['my_prediction_outputs'] tensor_info:
        dtype: DT_FLOAT
        shape: (-1, 1)
        name: StatefulPartitionedCall:0
  Method name is: tensorflow/serving/predict

- Eric McLachlan

---

感谢您的回答！我将尝试使用tf.Module进行测试，应该是类似的。 - Miguel Angel Alba Acosta

这里的 self.serving._model 是什么？它是否因为我们使用的装饰器而始终可用？ - Jash Shah

嗨@JashShah。很抱歉，self.serving._model是从本地项目复制粘贴留下来的。我已经删除了这个引用，并用model替换了它。所以，model只是TensorFlow对象。它在调用model.fit()时使用的是同一个对象。 - Eric McLachlan

Answer 2

创建 serving_default 签名的另一种方法是：

import tensorflow as tf
import tensorflow_hub as hub
import tensorflow_text

export_dir = "./models/use/00000001"
module = hub.load("https://tfhub.dev/google/universal-sentence-encoder-multilingual-large/3")

@tf.function
def my_module_encoder(text):
   inputs = {
        'text': text,
   }
   outputs = {
        'embeddings': module(text)
   }
   return outputs

tf.saved_model.save(
    module, 
    export_dir, 
    signatures=my_module_encoder.get_concrete_function(
        text=tf.TensorSpec(shape=None, dtype=tf.string)
    ), 
    options=None
)

你可以使用以下命令查看已创建的 `SignatureDefs` 签名：saved_model_cli。

$ saved_model_cli show --all  --dir models/use/00000001

MetaGraphDef with tag-set: 'serve' contains the following SignatureDefs:

signature_def['__saved_model_init_op']:
  The given SavedModel SignatureDef contains the following input(s):
  The given SavedModel SignatureDef contains the following output(s):
    outputs['__saved_model_init_op'] tensor_info:
        dtype: DT_INVALID
        shape: unknown_rank
        name: NoOp
  Method name is:

signature_def['serving_default']:
  The given SavedModel SignatureDef contains the following input(s):
    inputs['text'] tensor_info:
        dtype: DT_STRING
        shape: unknown_rank
        name: serving_default_text:0
  The given SavedModel SignatureDef contains the following output(s):
    outputs['embeddings'] tensor_info:
        dtype: DT_FLOAT
        shape: (-1, 512)
        name: StatefulPartitionedCall:0
  Method name is: tensorflow/serving/predict

- Hakan Özler

---

假设我的保存模型是无监督（聚类）模型，那么我应该使用哪种签名格式。正如他们在这个文档https://www.tensorflow.org/tfx/serving/signature_defs中提到的，对于无监督模型没有特定的格式。如果您有任何想法，请告诉我。 - user11530462