Tensorflow 2.0 具体函数 structured_input_signature return 值

Tensorflow 2.0 concrete function structured_input_signature return value

在检查 tf.ConcreteFunction 时,我很难理解 structured_input_signature 的 return 类型。

在 google 文档 https://www.tensorflow.org/guide/concrete_function#using_a_concrete_function 中,return 编辑了一个元组。例如

@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 的 return 元组中填充元组与字典的逻辑是什么?

简答

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 {
    }
  }
}

但是,上面的序列化结构并不能满足需求。我们无法将输入分配给键,因为 return 是一个元组。

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

您可能已经意识到,序列化模型的实际过程要复杂得多,其中涉及为服务添加新标签和签名、为分发策略添加副本内和交叉复制上下文,等等。不管复杂度如何,核心都是一样的,获取签名并序列化,代码来源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')})