Tensorflow 2.0的具体功能结构化输入签名返回值

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的返回元组中，为什么要选择填充元组或字典？这背后的逻辑是什么？

- applecider

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

- GabrielChu