tf.py_function

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

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

Используется в ноутбуках

Использование tf.py_function внутри tf.function позволяет запускать функцию Python с использованием немедленного выполнения внутри графа tf.function. Это имеет два основных эффекта:

1. Это позволяет использовать nofunc=None, inp=None, Tout=None код tensorflow внутри вашего tf.function.
2. Это позволяет запускать логику управления Python в tf.function без полагания на tf.autograph для преобразования кода в использование логики управления TensorFlow (tf.cond, tf.while_loop).

Обе эти функции могут быть полезны для отладки.

Поскольку tf.py_function работает с Tensor, она по-прежнему дифференцируема (один раз).

Есть два способа использования этой функции:

В качестве декоратора

Используйте tf.py_function в качестве декоратора, чтобы убедиться, что функция всегда запускается немедленно.

При использовании tf.py_function в качестве декоратора:

• вы должны установить Tout
• вы можете установить name
• вы не должны устанавливать func или inp

Например, вы можете использовать tf.py_function для реализации функции log huber.

@tf.py_function(Tout=tf.float32)
def py_log_huber(x, m):
  print('Running with eager execution.')
  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.constant(1.0)
m = tf.constant(2.0)

print(py_log_huber(x,m).numpy())
Running with eager execution.
1.0

Внутри tf.function tf.py_function не преобразуется в tf.Graph.:

@tf.function
def tf_wrapper(x):
  print('Tracing.')
  m = tf.constant(2.0)
  return py_log_huber(x,m)

Функция tf.py_function выполняется только немедленно и только при вызове tf.function:

print(tf_wrapper(x).numpy())
Tracing.
Running with eager execution.
1.0
print(tf_wrapper(x).numpy())
Running with eager execution.
1.0

Градиенты работают как ожидается:

with tf.GradientTape() as t:
  t.watch(x)
  y = tf_wrapper(x)
Running with eager execution.

t.gradient(y, x).numpy()
2.0

Помещенное

Вы также можете пропустить декоратор и использовать tf.py_function непосредственно. Эта форма является полезной сокращенной записью, если вы не контролируете исходный код функции, но её сложнее читать.

# No decorator
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.constant(1.0)
m = tf.constant(2.0)

tf.py_function(func=log_huber, inp=[x, m], Tout=tf.float32).numpy()
1.0

Дополнительная информация

Вы также можете использовать tf.py_function для отладки ваших моделей во время выполнения с использованием инструментов Python, т. е. вы можете изолировать части своего кода, которые вы хотите отлаживать, обернуть их в функции Python и вставить pdb точки отслеживания или операторы печати по желанию и обернуть эти функции в tf.py_function.

Дополнительную информацию о немедленном выполнении можно найти в Руководстве по немедленному выполнению.

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

Примечание: Мы рекомендуем избегать использования tf.py_function за пределами прототипирования и экспериментов из-за следующих известных ограничений:

• Вызов tf.py_function захватывает блокировку глобального интерпретатора Python (GIL), которая позволяет только одной нити выполняться в любой момент времени. Это предотвратит эффективную распараллеливание и распределение выполнения программы.

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

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

• В настоящее время tf.py_function несовместим с XLA. Вызов tf.py_function внутри tf.function(jit_compile=True) вызовет ошибку.