tf.py_function

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

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

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

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

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

Обе эти возможности могут быть полезны при отладке.

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

Эту функцию можно использовать двумя способами:

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

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

При использовании 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))

В eager-режиме функция работает обычным образом:

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 выполняется только в eager-режиме и только при вызове 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 или инструкции print, а затем обернуть эти функции в tf.py_function.

Дополнительные сведения о выполнении в eager-режиме см. в руководстве по eager-режиму.

tf.py_function по своей сути похожа на tf.numpy_function, но, в отличие от последней, позволяет использовать операции TensorFlow в обёрнутой функции Python. В частности, tf.compat.v1.py_func выполняется только на CPU и оборачивает функции, принимающие массивы 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) приведёт к ошибке.