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.