tf.numpy_function

Оборачивает функцию Python и использует ее как операцию TensorFlow.

tf.numpy_function(
    func=None, inp=None, Tout=None, stateful=True, name=None
)

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

Если задана функция Python func, оберните ее в операцию tf.function. func должна принимать массивы NumPy в качестве аргументов и возвращать массивы NumPy в качестве результатов.

Использовать tf.numpy_function можно двумя способами.

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

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

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

@tf.numpy_function(Tout=tf.float32)
def my_numpy_func(x):
  # x will be a numpy array with the contents of the input to the
  # tf.function
  print(f'executing eagerly, {x=}')
  return np.sinh(x)

Функция выполняется в eager-режиме:

my_numpy_func(1.0).numpy()
executing eagerly, x=1.0
1.17520

Поведение не меняется внутри tf.function:

@tf.function(input_signature=[tf.TensorSpec(None, tf.float32)])
def tf_function(input):
  y = tf.numpy_function(my_numpy_func, [input], tf.float32)
  return y
tf_function(tf.constant(1.)).numpy()
executing eagerly, x=array(1.)
1.17520

Вызов без декоратора

Этот вариант может быть полезен, если вы не контролируете исходный код функции, но его труднее читать.

Ниже приведена та же функция без декоратора:

def my_func(x):
  # x will be a numpy array with the contents of the input to the
  # tf.function
  print(f'executing eagerly, {x=}')
  return np.sinh(x)

Чтобы вызвать tf.numpy_function напрямую, передайте функцию, ее входные данные и тип выходных данных одним вызовом tf.numpy_function:

tf.numpy_function(my_func, [tf.constant(1.0)], tf.float32)
executing eagerly, x=array(1.)
1.17520

Подробнее

Сравнение с tf.py_function: tf.py_function и tf.numpy_function очень похожи, за исключением того, что tf.numpy_function принимает массивы NumPy, а не tf.Tensorы. Если функция должна содержать tf.Tensors, а все операции TensorFlow, выполняемые в функции, должны поддерживать дифференцирование, используйте tf.py_function.

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

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

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

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

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

• Поскольку функция принимает массивы NumPy, нельзя вычислить градиенты через numpy_function. Если вам нужна возможность дифференцирования, рассмотрите возможность использования tf.py_function.