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

至此模型创建正确，签名可以正确导出

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

Tensorflow saved model does not contain input names

python

tensorflow

keras

tensorflow-serving

tensorflow2.x