Wraps a python function into a TensorFlow op that executes it eagerly.
Aliases:
tf.compat.v1.py_function
tf.compat.v2.py_function
tf.py_function(
func,
inp,
Tout,
name=None
)
Used in the guide:
t
f.
da
ta
:B
u
il
dT
e
ns
or
Fl
ow
i
np
ut
p
ip
el
in
es
Used in the tutorials:
B
et
te
rp
e
rf
or
ma
nc
ew
i
th
t
f.
fu
nc
ti
on
L
oa
dt
e
xt
T
FR
ec
or
da
n
dt
f
.E
xa
mp
le
T
ra
ns
fo
rm
er
m
od
el
f
or
l
an
gu
ag
eu
n
de
rs
ta
nd
in
g`` tf.py_functionThis function allows expressing computations in a TensorFlow graph as Python functions. In particular, it wraps a Python function func in a once-differentiable TensorFlow operation that executes it with eager execution enabled. As a consequence, makes it possible to express control flow using Python constructs (if, while, for, etc.), instead of TensorFlow control flow constructs (tf.cond, tf.while_loop). For example, you might use to implement the log huber function:
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_functionYou can also use to debug your models at runtime using Python tools, i.e., you can isolate portions of your code that you want to debug, wrap them in Python functions and insert pdb tracepoints or print statements as desired, and wrap those functions in .
Eager guideFor more information on eager execution, see the .
tf.py_function is similar in spirit to tf.compat.v1.py_func, but unlike the latter, the former lets you use TensorFlow operations in the wrapped Python function. In particular, while tf.compat.v1.py_func only runs on CPUs and wraps functions that take NumPy arrays as inputs and return NumPy arrays as outputs, can be placed on GPUs and wraps functions that take Tensors as inputs, execute TensorFlow operations in their bodies, and return Tensors as outputs.
tf.compat.v1.py_funcLike , tf.py_function has the following limitations with respect to serialization and distribution:
- The body of the
func
tion (i.e.func
) will not be serialized in aGraphDef
. Therefore, you should not use thisfunc
tion if you need to serialize your model and restore it in a different environment. - The operation must run in the same address space as the Python program that calls
tf.py_function
(). If you are using distributed TensorFlow, you must run atf.distribute.Server
in the same process as the program that callstf.py_function
() and you must pin the created operation to a device in that server (e.g. using with tf.device()😃.
Args:
func
: A Pythonfunc
tion which accepts a list ofTensor
objects having element types that match the correspondingtf.Tensor
objects ininp
and returns a list ofTensor
objects (or a singleTensor
, orNone
) having element types that match the corresponding values inTout
.inp
: A list ofTensor
objects.Tout
: A list or tuple of tensorflow data types or a single tensorflow data type if there is only one, indicating whatfunc
returns; an empty list if no value is returned (i.e., if the return value isNone
).name
: Aname
for the operation (optional).
Returns:
A list of Tensor
or a single Tensor
which func
computes; an empty list if func
returns None.