tf.numpy_function

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

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

Принимая функцию Python func обернуть эту функцию как операцию в функции TensorFlow. func должна принимать массивы NumPy в качестве аргументов и возвращать массивы NumPy в качестве выходов.

Следующий пример создаёт граф TensorFlow с np.sinh() как операцией в графе:

def my_numpy_func(x):
  # x will be a numpy array with the contents of the input to the
  # tf.function
  return np.sinh(x)
@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 * y
tf_function(tf.constant(1.))
<tf.Tensor: shape=(), dtype=float32, numpy=1.3810978>

Сравнение с 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_comiple=True) вызовет ошибку.

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