Tensorflow 2.0具体函数structural_input_signature返回值

Question

structured_input_signature在检查tf.ConcreteFunction.

在谷歌文档https://www.tensorflow.org/guide/concrete_function#using_a_concrete_function中返回一个元组。例如

@tf.function
def power(a,b):
  print('Tracing "power"\n')
  return a**b

float_power = power.get_concrete_function(
  a = tf.TensorSpec(shape=[], dtype=tf.float32),
  b = tf.TensorSpec(shape=[], dtype=tf.float32))

print(float_power.structured_input_signature)
print(float_power.structured_outputs)

印刷

Tracing "power"

((TensorSpec(shape=(), dtype=tf.float32, name='a'), TensorSpec(shape=(), dtype=tf.float32, name='b')), {})
Tensor("Identity:0", shape=(), dtype=float32)

然而，当模块被保存和加载时，输出略有不同：

float_power_mod = tf.Module()
float_power_mod.float_power = float_power
tf.saved_model.save(float_power_mod, './float_power_mod')

mod_4 = tf.saved_model.load('./float_power_mod')
float_power_func = mod_4.signatures['serving_default']
print(float_power_func.structured_input_signature)

印刷

((),
 {'a': TensorSpec(shape=(), dtype=tf.float32, name='a'),
  'b': TensorSpec(shape=(), dtype=tf.float32, name='b')})

在 Structured_input_signature 的返回元组中填充元组与字典背后的逻辑是什么？

Answer 1

简短回答

dict允许我们将关键字参数传递给函数，以便我们可以将实值输入张量标记到 TF 接受的相应占位符。

result = float_power_func(a=tf.constant(2.), b=tf.constant(3.))

长答案

为了保存 TF 模型，首先我们需要序列化张量。在导出的目录下，您可以找到一个.pb文件，即用于序列化整个模型的 protobuf。我所说的模型是指张量的集合以及这些张量之间的关系，所有这些都被捕获在 protobuf 中。而TF已经提供了序列化的功能，以你的代码为例

from tensorflow.python.saved_model import nested_structure_coder

coder = nested_structure_coder.StructureCoder()
signature_proto = coder.encode_structure(float_power.structured_input_signature)
print(signature_proto)

印刷

tuple_value {
  values {
    tuple_value {
      values {
        tensor_spec_value {
          name: "a"
          shape {
          }
          dtype: DT_FLOAT
        }
      }
      values {
        tensor_spec_value {
          name: "b"
          shape {
          }
          dtype: DT_FLOAT
        }
      }
    }
  }
  values {
    dict_value {
    }
  }
}

然而，上述序列化结构并不能满足需要。我们无法将输入分配给键，因为返回的是一个元组。

((TensorSpec(shape=(), dtype=tf.float32, name='a'), TensorSpec(shape=(), dtype=tf.float32, name='b')), {})

正如您可能意识到的，序列化模型的实际过程要复杂得多，其中涉及添加用于服务的新标签和签名、用于分发策略的副本内和跨副本上下文等。无论复杂程度如何，核心都是相同的，获取签名并将其序列化，代码来源于此处

signatures = signature_serialization.canonicalize_signatures(signatures)

重新打包并将输入张量作为键值signatures对移动到内部dict_value

value {
    canonicalized_input_signature {
      tuple_value {
        values {
          tuple_value {
          }
        }
        values {
          dict_value {
            fields {
              key: "a"
              value {
                tensor_spec_value {
                  name: "a"
                  shape {
                  }
                  dtype: DT_FLOAT
                }
              }
            }
            fields {
              key: "b"
              value {
                tensor_spec_value {
                  name: "b"
                  shape {
                  }
                  dtype: DT_FLOAT
                }
              }
            }
          }
        }
      }
    }

并解码你会得到

((),
 {'a': TensorSpec(shape=(), dtype=tf.float32, name='a'),
  'b': TensorSpec(shape=(), dtype=tf.float32, name='b')})