python - Tensorflow 保存的模型不包含输入名称

Question

我们目前正在 tensorflow 2.4.0 中训练一个对象检测模型，该模型运行良好。然而，为了能够为它服务，我们需要用一个图像预处理层包装它，该层将图像字节作为输入并将它们转换为检测模型所需的图像张量。请参阅以下代码:

png_file = 'myfile.png'
input_tensor = tf.io.read_file(png_file, name='image_bytes')

def preprocessing_layer(inputs):

    image_tensor = tf.image.decode_image(inputs, channels=3)
    image_tensor = tf.expand_dims(
        image_tensor, axis=0, name=None
    )
    return image_tensor 

model = keras.Sequential(
    [
        keras.Input(tensor=input_tensor, dtype=tf.dtypes.string, name='image_bytes', batch_size=1),
        tf.keras.layers.Lambda(lambda inp: preprocessing_layer(inp)),
        yolo_model
    ]
)
model.summary()

这个包装模型提供了有用的检测，如果我们调用 model.input_names返回正确的名称: ['image_bytes'] .
现在，如果我们使用 model.save('model_path') 保存模型保存的模型不再包含输入名称，而是将它们替换为通用名称 ( args_0 )。

signature_def['serving_default']:
  The given SavedModel SignatureDef contains the following input(s):
    inputs['args_0'] tensor_info:
        dtype: DT_STRING
        shape: ()
        name: serving_default_args_0:0
  The given SavedModel SignatureDef contains the following output(s):
    outputs['model'] tensor_info:
        dtype: DT_FLOAT
        shape: (1, 64512, 6)

这是一个问题，因为 tensorflow 服务依赖于以 _bytes 结尾的名称。转换 base64 输入。
您能否提供有关在保存模型时如何保留输入名称的提示？

Answer 1

问题源于您定义 lambda 层的方式以及您设置模型的方式。
您的 lambda 函数应该能够处理批处理，目前情况并非如此。你可以天真地使用 tf.map_fn使其处理一批图像，如下所示:

def preprocessing_layer(str_inputs):
    def decode(inputs):
        image_tensor = tf.image.decode_image(inputs[0], channels=3)
        image_tensor = tf.expand_dims(
            image_tensor, axis=0, name=None
        )
        return image_tensor
    return tf.map_fn(decode, str_inputs, fn_output_signature=tf.uint8)

然后您可以使用符号 tf.keras.Input 定义您的模型。 ，将形状设置为 () (指定除批量大小之外的其他维度):

model = keras.Sequential(
    [
        keras.Input((), dtype=tf.dtypes.string, name='image_bytes'),
        tf.keras.layers.Lambda(lambda inp: preprocessing_layer(inp)),
        yolo_model
    ]
)

现在模型已正确创建，签名可以正确导出。