Tensorflow 2.0 具体函数structured_input_signature 返回值

Question

在检查tf.ConcreteFunction 时，我很难理解structured_input_signature 的返回类型。

在谷歌文档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')), {})

您可能已经意识到，序列化模型的实际过程要复杂得多，其中包括为服务添加新标签和签名、为分发策略添加副本内和跨副本上下文等。 Regradless所有的复杂性，核心都是一样的，获取签名并序列化，代码来源here

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')})