此页面由 Cloud Translation API 翻译。
Switch to English

tf.py_function

TensorFlow 1版 GitHub上查看源代码

包装一个Python函数为TensorFlow运算是急切地执行它。

用在笔记本电脑

使用的指南使用教程

该功能可以以一个TensorFlow图作为Python函数表达计算。特别是,它包装一个Python函数func与渴望执行启用执行它一次可微TensorFlow操作。因此, tf.py_function使得能够使用Python构建体(表达控制流程ifwhilefor等),而不是TensorFlow控制流结构( tf.condtf.while_loop )。例如,你可以使用tf.py_function实现日志HUBER功能:

 def log_huber(x, m):
  if tf.abs(x) <= m:
    return x**2
  else:
    return m**2 * (1 - 2 * tf.math.log(m) + tf.math.log(x**2))

x = tf.compat.v1.placeholder(tf.float32)
m = tf.compat.v1.placeholder(tf.float32)

y = tf.py_function(func=log_huber, inp=[x, m], Tout=tf.float32)
dy_dx = tf.gradients(y, x)[0]

with tf.compat.v1.Session() as sess:
  # The session executes `log_huber` eagerly. Given the feed values below,
  # it will take the first branch, so `y` evaluates to 1.0 and
  # `dy_dx` evaluates to 2.0.
  y, dy_dx = sess.run([y, dy_dx], feed_dict={x: 1.0, m: 2.0})
 

您还可以使用tf.py_function在使用Python的工具,也就是说,您可以隔离您要调试的代码部分的运行调试你的模型,在Python功能和插入他们包裹pdb跟踪点,或根据需要打印报表,和那些包裹在功能tf.py_function

有关急于执行的详细信息,请参阅渴望指南

tf.py_function在本质上类似tf.compat.v1.py_func ,但不同的是后者,前者可以让你在包装的Python函数使用TensorFlow操作。特别是,虽然tf.compat.v1.py_func只能运行在CPU和包装了采取NumPy的阵列作为输入,并返回NumPy的阵列为输出函数, tf.py_function可以被放置在该采取张量作为输入GPU和包裹物的功能,执行TensorFlow在他们的身体运作,并返回张量为输出。

tf.compat.v1.py_functf.py_function相对于系列化和经销以下限制:

  • 的函数的主体(即func )将不会在被序列化GraphDef 。因此,如果您需要序列化你的模型,它在不同的环境恢复你不应该使用此功能。

  • 操作必须为Python程序相同的地址空间调用运行tf.py_function()如果你使用分布式TensorFlow,你必须运行tf.distribute.Server在同一进程中调用该程序tf.py_function()您必须在引脚该服务器的创建操作的设备(例如,使用with tf.device():

func 它接受的列表的Python函数Tensor对象具有匹配相应的那个元素类型tf.Tensor在对象inp和返回一个列表的Tensor对象(或单个Tensor ,或None具有匹配在对应的值即元素类型) Tout
inp 清单Tensor的对象。
Tout tensorflow数据类型或单一tensorflow数据类型的列表或元组,如果仅存在一个,指示什么func返回;一个空列表,如果没有返回值是(即,如果返回值是None )。
name 一种操作(可选)的名称。

的列表Tensor或单个Tensor ,其func单位计算;一个空列表,如果func返回无。