tf.py_function

Оборачивает функцию Python в операцию TensorFlow, которая выполняет её немедленно.

tf.py_function(
    func, inp, Tout, name=None
)

Эта функция позволяет выражать вычисления в графике TensorFlow в виде функций Python. В частности, она оборачивает функцию Python func в однократно дифференцируемую операцию TensorFlow, которая выполняет её с включённым немедленным выполнением. В результате, tf.py_function позволяет выразить поток управления с помощью конструкций Python (if, while, for, и т. д.), вместо конструкций потока управления TensorFlow (tf.cond, tf.while_loop). Например, вы можете использовать tf.py_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_function для отладки ваших моделей во время выполнения с помощью инструментов Python, т. е. вы можете изолировать части вашего кода, которые вы хотите отладить, обернуть их в функции Python и вставить pdb точки отслеживания или операторы вывода, как требуется, и обернуть эти функции в tf.py_function.

Для получения более подробной информации о немедленном выполнении, см. Руководство по немедленному выполнению.

tf.py_function похожа по духу на tf.compat.v1.py_func, но в отличие от последней, первая позволяет использовать операции TensorFlow в обернутой функции Python. В частности, в то время как tf.compat.v1.py_func работает только на ЦП и оборачивает функции, принимающие массивы NumPy в качестве входных данных и возвращающие массивы NumPy в качестве выходных данных, tf.py_function может размещаться на GPU и оборачивает функции, принимающие тензоры в качестве входных данных, выполняющие операции TensorFlow в своих телах и возвращающие тензоры в качестве выходных данных.

Подобно tf.compat.v1.py_func, tf.py_function имеет следующие ограничения в отношении сериализации и распределения:

• Тело функции (т. е. func) не будет сериализовано в GraphDef. Поэтому не следует использовать эту функцию, если вам необходимо сериализовать вашу модель и восстановить её в другой среде.

• Операция должна выполняться в том же адресном пространстве, что и программа Python, которая вызывает tf.py_function(). Если вы используете распределённый TensorFlow, вы должны запустить tf.distribute.Server в том же процессе, что и программа, вызывающая tf.py_function(), и вы должны привязать созданную операцию к устройству в этом сервере (например, с помощью with tf.device():).