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

问题描述

我想更改保存的模型的输入和输出签名，我使用 tf.Module 对象来构建主模型的操作.

I would like to change the input and output signatures of the model saved, I used tf.Module objects to build the operations of the main model.

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"作为签名，我该如何更改?

then I am loading the model but I have as signatures 'default_serving' and 'output_0', how can I change this?

推荐答案

我想出了一种不使用 tf.Module 定义输出签名的方法，方法是定义一个 tf.function 返回一个字典字典中使用的键将作为输出名称的输出.

I figured out a way to define the output signature without using tf.Module by defining a tf.function that returns a dictionary of outputs where the keys used in the dictionary will be the output names.

# 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

这篇关于Tensorflow 2.0:如何在使用 tf.saved_model 时更改输出签名的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持！