python - 从 python API 而不是 saved_model_cli 中提取(或设置)输入/输出 TF 张量名称信息

Question

我用 Keras/TF2.5 训练了一个简单的模型并将其保存为已保存的模型。

tf.saved_model.save(my_model,'/path/to/model')

如果我通过

检查它

saved_model_cli show --dir /path/to/model --tag_set serve --signature_def serving_default

我得到这些输出/名称:

inputs['conv2d_input'] tensor_info:
  dtype: DT_FLOAT
  shape: (-1, 32, 32, 1)
  name: serving_default_conv2d_input:0
outputs['dense'] tensor_info:
  dtype: DT_FLOAT
  shape: (-1, 2)
  name: StatefulPartitionedCall:0

名称 serving_default_conv2d_input 和 StatefulPartitionedCall 实际上可以用于推理。

我想使用 python API 提取它们。如果我通过加载模型来查询它:

>>> m=tf.saved_model.load('/path/to/model')
>>> m.signatures['serving_default'].inputs[0].name
'conv2d_input:0'
>>> m.signatures['serving_default'].outputs[0].name
'Identity:0'

我得到完全不同的名字。

问题:

1. 如何从 python API 中提取这些名称 serving_default_conv2d_input 和 StatefulPartitionedCall？
2. 或者，当我调用 tf.saved_model.save 时如何定义/修复名称？
3. :0 是什么意思？

附带问题:

您如何通过 SavedModel 将 TF 模型部署到生产环境？

Answer 1

saved_model_cli显示的输入/输出张量名称可以提取如下:

from tensorflow.python.tools import saved_model_utils

saved_model_dir = '/path/to/model'
tag_set = 'serve'
signature_def_key = 'serving_default'

# 1. Load MetaGraphDef with saved_model_utils
meta_graph_def = saved_model_utils.get_meta_graph_def(saved_model_dir, tag_set)

# 2. Get input signature names
input_signatures = list(meta_graph_def.signature_def[signature_def_key].inputs.values())
input_names = [signature.name for signature in input_signatures]
print(input_names)  # ['serving_default_conv2d_input:0']

# 3. Get output signature names
output_signatures = list(meta_graph_def.signature_def[signature_def_key].outputs.values())
output_names = [signature.name for signature in output_signatures]
print(output_names)  # ['StatefulPartitionedCall:0']

关于 :0 的含义，op_name:0 表示“作为名为 op_name 的操作的第 0 个输出的张量。”因此，您可能会使用 …:1 来获取具有多个输出的操作的输出，但许多操作都是单输出的，因此您将始终对它们使用 …:0 (来源:@mrry's comment)。