tf.py_function
Stay organized with collections
Save and categorize content based on your preferences.
Wraps a python function into a TensorFlow op that executes it eagerly.
tf.py_function(
func, inp, Tout, name=None
)
This 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, tf.py_function
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 tf.py_function
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})
You can also use tf.py_function
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
tf.py_function
.
For more information on eager execution, see the
Eager guide.
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, tf.py_function
can be placed on GPUs and wraps functions
that take Tensors as inputs, execute TensorFlow operations in their bodies,
and return Tensors as outputs.
Like tf.compat.v1.py_func
, tf.py_function
has the following limitations
with respect to serialization and distribution:
The body of the function (i.e. func
) will not be serialized in a
GraphDef
. Therefore, you should not use this function 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 a tf.distribute.Server
in the same process as the
program that calls tf.py_function()
and you must pin the created
operation to a device in that server (e.g. using with tf.device():
).
Args |
func
|
A Python function that accepts inp as arguments, and returns a
value (or list of values) whose type is described by Tout .
|
inp
|
Input arguments for func . A list whose elements are Tensor s or
CompositeTensors (such as tf.RaggedTensor ); or a single Tensor or
CompositeTensor .
|
Tout
|
The type(s) of the value(s) returned by func . One of the
following.
If func returns a Tensor (or a value that can be converted to a
Tensor): the tf.DType for that value.
If func returns a CompositeTensor : The tf.TypeSpec for that value.
If func returns None : the empty list ([] ).
If func returns a list of Tensor and CompositeTensor values:
a corresponding list of tf.DType s and tf.TypeSpec s for each value.
|
name
|
A name for the operation (optional).
|
Returns |
The value(s) computed by func : a Tensor , CompositeTensor , or list of
Tensor and CompositeTensor ; or an empty list if func returns None .
|
Except as otherwise noted, the content of this page is licensed under the Creative Commons Attribution 4.0 License, and code samples are licensed under the Apache 2.0 License. For details, see the Google Developers Site Policies. Java is a registered trademark of Oracle and/or its affiliates. Some content is licensed under the numpy license.
Last updated 2022-11-04 UTC.
[[["Easy to understand","easyToUnderstand","thumb-up"],["Solved my problem","solvedMyProblem","thumb-up"],["Other","otherUp","thumb-up"]],[["Missing the information I need","missingTheInformationINeed","thumb-down"],["Too complicated / too many steps","tooComplicatedTooManySteps","thumb-down"],["Out of date","outOfDate","thumb-down"],["Samples / code issue","samplesCodeIssue","thumb-down"],["Other","otherDown","thumb-down"]],["Last updated 2022-11-04 UTC."],[],[],null,["# tf.py_function\n\n\u003cbr /\u003e\n\n|-----------------------------------------------------------------------------------------------------------------------------|\n| [View source on GitHub](https://github.com/tensorflow/tensorflow/blob/v2.7.4/tensorflow/python/ops/script_ops.py#L427-L523) |\n\nWraps a python function into a TensorFlow op that executes it eagerly.\n\n#### View aliases\n\n\n**Compat aliases for migration**\n\nSee\n[Migration guide](https://www.tensorflow.org/guide/migrate) for\nmore details.\n\n[`tf.compat.v1.py_function`](https://www.tensorflow.org/api_docs/python/tf/py_function)\n\n\u003cbr /\u003e\n\n tf.py_function(\n func, inp, Tout, name=None\n )\n\nThis function allows expressing computations in a TensorFlow graph as\nPython functions. In particular, it wraps a Python function `func`\nin a once-differentiable TensorFlow operation that executes it with eager\nexecution enabled. As a consequence, [`tf.py_function`](../tf/py_function) makes it\npossible to express control flow using Python constructs (`if`, `while`,\n`for`, etc.), instead of TensorFlow control flow constructs ([`tf.cond`](../tf/cond),\n[`tf.while_loop`](../tf/while_loop)). For example, you might use [`tf.py_function`](../tf/py_function) to\nimplement the log huber function: \n\n def log_huber(x, m):\n if tf.abs(x) \u003c= m:\n return x**2\n else:\n return m**2 * (1 - 2 * tf.math.log(m) + tf.math.log(x**2))\n\n x = tf.compat.v1.placeholder(tf.float32)\n m = tf.compat.v1.placeholder(tf.float32)\n\n y = tf.py_function(func=log_huber, inp=[x, m], Tout=tf.float32)\n dy_dx = tf.gradients(y, x)[0]\n\n with tf.compat.v1.Session() as sess:\n # The session executes `log_huber` eagerly. Given the feed values below,\n # it will take the first branch, so `y` evaluates to 1.0 and\n # `dy_dx` evaluates to 2.0.\n y, dy_dx = sess.run([y, dy_dx], feed_dict={x: 1.0, m: 2.0})\n\nYou can also use [`tf.py_function`](../tf/py_function) to debug your models at runtime\nusing Python tools, i.e., you can isolate portions of your code that\nyou want to debug, wrap them in Python functions and insert `pdb` tracepoints\nor print statements as desired, and wrap those functions in\n[`tf.py_function`](../tf/py_function).\n\nFor more information on eager execution, see the\n[Eager guide](https://tensorflow.org/guide/eager).\n\n[`tf.py_function`](../tf/py_function) is similar in spirit to [`tf.compat.v1.py_func`](../tf/compat/v1/py_func), but unlike\nthe latter, the former lets you use TensorFlow operations in the wrapped\nPython function. In particular, while [`tf.compat.v1.py_func`](../tf/compat/v1/py_func) only runs on CPUs\nand\nwraps functions that take NumPy arrays as inputs and return NumPy arrays as\noutputs, [`tf.py_function`](../tf/py_function) can be placed on GPUs and wraps functions\nthat take Tensors as inputs, execute TensorFlow operations in their bodies,\nand return Tensors as outputs.\n\nLike [`tf.compat.v1.py_func`](../tf/compat/v1/py_func), [`tf.py_function`](../tf/py_function) has the following limitations\nwith respect to serialization and distribution:\n\n- The body of the function (i.e. `func`) will not be serialized in a\n `GraphDef`. Therefore, you should not use this function if you need to\n serialize your model and restore it in a different environment.\n\n- The operation must run in the same address space as the Python program\n that calls [`tf.py_function()`](../tf/py_function). If you are using distributed\n TensorFlow, you must run a [`tf.distribute.Server`](../tf/distribute/Server) in the same process as the\n program that calls [`tf.py_function()`](../tf/py_function) and you must pin the created\n operation to a device in that server (e.g. using `with tf.device():`).\n\n\u003cbr /\u003e\n\n\u003cbr /\u003e\n\n\u003cbr /\u003e\n\n| Args ---- ||\n|--------|------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|\n| `func` | A Python function that accepts `inp` as arguments, and returns a value (or list of values) whose type is described by `Tout`. |\n| `inp` | Input arguments for `func`. A list whose elements are `Tensor`s or `CompositeTensors` (such as [`tf.RaggedTensor`](../tf/RaggedTensor)); or a single `Tensor` or `CompositeTensor`. |\n| `Tout` | The type(s) of the value(s) returned by `func`. One of the following. \u003cbr /\u003e - If `func` returns a `Tensor` (or a value that can be converted to a Tensor): the [`tf.DType`](../tf/dtypes/DType) for that value. - If `func` returns a `CompositeTensor`: The [`tf.TypeSpec`](../tf/TypeSpec) for that value. - If `func` returns `None`: the empty list (`[]`). - If `func` returns a list of `Tensor` and `CompositeTensor` values: a corresponding list of [`tf.DType`](../tf/dtypes/DType)s and [`tf.TypeSpec`](../tf/TypeSpec)s for each value. |\n| `name` | A name for the operation (optional). |\n\n\u003cbr /\u003e\n\n\u003cbr /\u003e\n\n\u003cbr /\u003e\n\n\u003cbr /\u003e\n\n| Returns ------- ||\n|---|---|\n| The value(s) computed by `func`: a `Tensor`, `CompositeTensor`, or list of `Tensor` and `CompositeTensor`; or an empty list if `func` returns `None`. ||\n\n\u003cbr /\u003e"]]