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

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