TensorFlow 2.0 函数签名与图结构解析

陈默
陈默
发布时间:2026-10-13 | 5人浏览过
原创

本文深入解析TensorFlow 2.0中input_signature的作用及其与图结构的关系。通过具体代码示例,展示如何利用input_signature限定输入类型以增强函数安全性,并详细说明如何使用get_concrete_function将Python函数转换为可保存的ConcreteFunction对象,进而生成SavedModel。文章还涵盖了从ConcreteFunction中提取图定义、操作节点及张量名称的具体操作方法,旨在帮助开发者理解模型保存与加载的核心机制。

input_signature的作用与SavedModel保存机制

使用input_signature的主要优势在于对函数输入类型进行严格限制,从而有效防止在调用函数时因参数类型不匹配而引发错误。

首先,input_signature能够明确限定函数的输入数据类型,避免在调用过程中出现类型错误。

其次,在TensorFlow框架中,只有具备input_signature的函数才能被保存为SavedModel格式。在保存过程中,必须通过get_concrete_function函数,将带有@tf.function装饰器的普通Python函数转换为包含具体图定义的函数。

以下代码示例具体展示了input_signature如何限制函数的输入类型。


@tf.function(input_signature=[tf.TensorSpec([None], tf.int32, name='x')])
def cube(z): #实现输入的立方
 return tf.pow(z, 3)
try:
 print(cube(tf.constant([1., 2., 3.])))
except ValueError as ex:
 print(ex)
print(cube(tf.constant([1, 2, 3])))

运行结果如下:

Python inputs incompatible with input_signature: inputs: ( tf.Tensor([1. 2. 3.], shape=(3,), dtype=float32)) input_signature: ( TensorSpec(shape=(None,), dtype=tf.int32, name='x')) tf.Tensor([ 1 8 27], shape=(3,), dtype=int32)

get_concrete_function的使用与ConcreteFunction对象

关于get_concrete_function的具体用法

需要注意的是,下文介绍的函数在常规的模型构建和训练阶段通常不会直接使用。它们主要应用于两个场景:一是如何正确保存模型,二是保存后如何准确载入模型。

可以为由@tf.function装饰的普通Python函数添加input_signature,从而将该Python函数转换为可保存的TensorFlow图结构(即SavedModel)。

以下是函数用法的具体示例:


@tf.function(input_signature=[tf.TensorSpec([None], tf.int32, name='x')])
def cube(z):
 return tf.pow(z, 3)
 
try:
 print(cube(tf.constant([1., 2., 3.])))
except ValueError as ex:
 print(ex)
 
print(cube(tf.constant([1, 2, 3])))
 
# @tf.function py func -> tf graph
# get_concrete_function -> add input signature -> SavedModel
 
cube_func_int32 = cube.get_concrete_function(
 tf.TensorSpec([None], tf.int32)) #tensorflow的类型
print(cube_func_int32)

输出结果如下:

<tensorflow.python.eager.function.ConcreteFunction object at 0x00000240E29695C0>

从上述输出可以看出,调用get_concrete_function函数后,返回的是一个ConcreteFunction对象。


#看用新参数获得的对象与原来的对象是否一样
print(cube_func_int32 is cube.get_concrete_function(
 tf.TensorSpec([5], tf.int32))) #输入大小为5
print(cube_func_int32 is cube.get_concrete_function(
 tf.constant([1, 2, 3]))) #传具体数据

再次输出:

True True

从ConcreteFunction中提取图定义与节点信息

cube_func_int32.graph #获取图定义

输出结果如下:


[<tf.Operation 'x' type=Placeholder>,
 <tf.Operation 'Pow/y' type=Const>,
 <tf.Operation 'Pow' type=Pow>,
 <tf.Operation 'Identity' type=Identity>]

pow_op = cube_func_int32.graph.get_operations()[2]
print(pow_op)

输出内容如下:

name: "Pow" op: "Pow" input: "x" input: "Pow/y" attr { key: "T" value { type: DT_INT32 } }


print(list(pow_op.inputs))
print(list(pow_op.outputs))

输出结果如下:

[<tf.Tensor 'x:0' shape=(None,) dtype=int32>, <tf.Tensor 'Pow/y:0' shape=() dtype=int32>] [<tf.Tensor 'Pow:0' shape=(None,) dtype=int32>]

cube_func_int32.graph.get_operation_by_name("x")

输出结果如下:

<tf.Operation 'x' type=Placeholder>

cube_func_int32.graph.get_tensor_by_name("x:0") #注意默认需加上“:0”后缀

<tf.Tensor 'x:0' shape=(None,) dtype=int32>

cube_func_int32.graph.as_graph_def() #获取总图定义,适用于上述两种操作


node {
 name: "x"
 op: "Placeholder"
 attr {
 key: "_user_specified_name"
 value {
 s: "x"
 }
 }
 attr {
 key: "dtype"
 value {
 type: DT_INT32
 }
 }
 attr {
 key: "shape"
 value {
 shape {
 dim {
  size: -1
 }
 }
 }
 }
}
node {
 name: "Pow/y"
 op: "Const"
 attr {
 key: "dtype"
 value {
 type: DT_INT32
 }
 }
 attr {
 key: "value"
 value {
 tensor {
 dtype: DT_INT32
 tensor_shape {
 }
 int_val: 3
 }
 }
 }
}
node {
 name: "Pow"
 op: "Pow"
 input: "x"
 input: "Pow/y"
 attr {
 key: "T"
 value {
 type: DT_INT32
 }
 }
}
node {
 name: "Identity"
 op: "Identity"
 input: "Pow"
 attr {
 key: "T"
 value {
 type: DT_INT32
 }
 }
}
versions {
 producer: 119
}

至此,关于TensorFlow 2.0函数签名与图结构的介绍已完成。如需了解更多关于TensorFlow函数签名与图结构的内容,建议查阅本站过往文章或继续浏览后续相关文章,感谢大家的支持!

相关标签:

本站声明

本文内容仅供参考,不代表本站观点。转载或引用请注明出处。如发现内容问题,请通过联系我们页面反馈。

Python 相关产品和应用

全部下载

专题列表