Python tf.types.experimental.GenericFunction.get_concrete_function用法及代码示例

用法

get_concrete_function(
    *args, **kwargs
) -> tf.types.experimental.ConcreteFunction

返回专用于输入类型的 ConcreteFunction。

args 和kwargs 指定的参数遵循正常的函数调用规则。返回的 ConcreteFunction 具有与 self 相同的位置和关键字参数集，但它们的类型与 args 和 kwargs 指定的类型兼容(尽管不一定相等)。

@tf.function
def f(x):
  return x
f_concrete = f.get_concrete_function(tf.constant(1.0))
f_concrete = f.get_concrete_function(x=tf.constant(1.0))

与普通调用不同，get_concrete_function 允许使用类型说明符而不是 TensorFlow 对象，因此例如 tf.Tensor 可以替换为 tf.TensorSpec 。

@tf.function
def f(x):
  return x
f_concrete = f.get_concrete_function(tf.TensorSpec([], tf.float64))

如果函数定义只允许一种特化，args 和kwargs 可以完全省略。

@tf.function(input_signature=[tf.TensorSpec(None, tf.float32)])
def f(x):
  return x
f_concrete = f.get_concrete_function()

返回的ConcreteFunction可以正常调用：

f_concrete(tf.constant(1.0))
<tf.Tensor:shape=(), dtype=float32, numpy=1.0>
f_concrete(x=tf.constant(1.0))
<tf.Tensor:shape=(), dtype=float32, numpy=1.0>