tf.contrib.eager.defun

Компилирует функцию Python в вызываемый граф TensorFlow.

tf.contrib.eager.defun(
    func=None, input_signature=None, autograph=True,
    experimental_autograph_options=None, experimental_relax_shapes=False
)

defun (сокращённо "определить функцию") компилирует функцию Python, состоящую из операций TensorFlow, в вызываемый объект, который выполняет tf.Graph, содержащий эти операции. Вызываемый объект, созданный defun, содержит только подграф операций TensorFlow, которые были выполнены, когда функция Python была вызвана с определённой сигнатурой входных данных, определённой как список форм и типов аргументов функции Python, имеющих тип Tensor, и значений её не-Tensor аргументов Python.

При включённом режиме выполнения eager, возможность создавать графы из функций Python позволяет постепенно отказаться от отладки и интерактивности в пользу производительности. Функции, скомпилированные с помощью defun, нельзя просмотреть с помощью pdb; однако, выполнение графа, сгенерированного defun, иногда занимает меньше времени и памяти, чем выполнение соответствующей функции Python в режиме eager, так как представление вычислений в виде графов позволяет использовать оптимизации, такие как автоматическое повторное использование буферов и параллелизация между операциями. Обратите внимание, что выполнение функции, скомпилированной с помощью defun, влечёт за собой небольшую постоянную издержку, поэтому выполнение достаточно маленьких функций Python в режиме eager может занимать меньше времени, чем выполнение соответствующих графов, сгенерированных с помощью defun.

Для того, чтобы функция Python была совместима с defun, все её аргументы должны быть хешируемыми объектами Python или списками таких объектов. Функция не должна изменять структуру списка/словаря своих аргументов. Кроме того, она должна возвращать ноль или более tf.Tensor объектов. Если функция Python возвращает tf.Variable, её скомпилированная версия вернёт значение этой переменной в виде tf.Tensor.

Выполнение графа, сгенерированного defun, учитывает аннотации устройства (то есть все with tf.device директивы, присутствующие в функции Python, также будут присутствовать в соответствующем графе), но пока ещё невозможно выполнить сгенерированные графы на нескольких машинах.

Пример использования

import tensorflow as tf

tf.compat.v1.enable_eager_execution()

# A simple example.
def f(x, y):
  return tf.reduce_mean(tf.multiply(x ** 2, 3) + y)

g = tf.contrib.eager.defun(f)

x = tf.constant([[2.0, 3.0]])
y = tf.constant([[3.0, -2.0]])

# `f` and `g` will return the same value, but `g` will be executed as a
# TensorFlow graph.
assert f(x, y).numpy() == g(x, y).numpy()

# `defun` is capable of compiling Python functions that close over Python
# objects, including Tensors and Variables.
@tf.contrib.eager.defun
def h():
  return f(x, y)

assert (h().numpy() == f(x, y).numpy()).all()

# `defun` automatically lifts variables out of the graphs it creates,
# allowing you to compile the `call` methods of `tf.keras.layers.Layer` and
# `tf.keras.Model` objects.
class MyModel(tf.keras.Model):

  def __init__(self, keep_probability=0.2):
    super(MyModel, self).__init__()
    self.dense1 = tf.keras.layers.Dense(4, activation=tf.nn.relu)
    self.dense2 = tf.keras.layers.Dense(5, activation=tf.nn.softmax)
    self.keep_probability = keep_probability

  @tf.contrib.eager.defun
  def call(self, inputs, training=True):
    x = self.dense2(self.dense1(inputs))
    if training:
      return tf.nn.dropout(x, self.keep_probability)
    else:
      return x

model = MyModel()
model(x, training=True)  # executes a graph, with dropout
model(x, training=False) # executes a graph, without dropout

# `defun`-compiled functions are differentiable.
optimizer = tf.compat.v1.train.GradientDescentOptimizer(learning_rate=0.01)
with tf.GradientTape() as tape:
  outputs = model(x)
gradient = tape.gradient(outputs, model.trainable_variables)
optimizer.apply_gradients((grad, var) for grad, var in zip(gradient,
                          model.trainable_variables))

При использовании defun, есть тонкости, связанные с входными данными, потоком управления Python и созданием переменных, о которых необходимо знать. Для наглядности, пусть f — это функция Python, которая возвращает ноль или более tf.Tensor объектов, и пусть F = defun(f). F строит граф для каждой уникальной сигнатуры входных данных, поток управления Python встраивается в графы, и операции, связанные с инициализацией переменных, автоматически выносятся за пределы графов, которые F генерирует, и помещаются в контекст eager при выполнении в режиме eager или в внешний граф в противном случае.

Сигнатуры входных данных

По умолчанию, F = tf.contrib.eager.defun(f) создаёт отдельный граф для каждой уникальной последовательности форм и типов аргументов Tensor и значений объектов Python, с которыми она вызывается. Например, вызов F(tf.random.uniform([2]) выполнит другой граф, чем F(tf.random.uniform([3]), потому что два входных значения имеют разные формы. В первый раз, когда F(*args, **kwargs) вызывается с определённой последовательностью форм и типов Tensor и значений Python, она строит граф, прослеживая выполнение f(*args, **kwargs); этот граф связан с сигнатурой входных данных, выведенной из (*args, **kwargs), и кэшируется для последующего использования.

Массивы NumPy, передаваемые в качестве входных данных F, преобразуются в tf.Tensor объекты перед передачей в f, и обрабатываются как тензоры для кэширования. Это позволяет вызывать функцию несколько раз с массивами NumPy, имеющими разные значения, но одинаковую форму и тип, не пересчитывая каждый раз.

tf.contrib.eager.defun кэширует графы для удобства, позволяя определять функции TensorFlow без явного указания их сигнатур. Однако эта политика является консервативной и потенциально дорогостоящей; например, когда различные вызовы вашей функции имеют тензорные входные данные с разными формами, эта политика может генерировать больше функций графов, чем необходимо. Чтобы устранить такие затраты, tf.contrib.eager.defun позволяет указать необязательный аргумент input_signature, определяющий формы и типы входных данных. В частности, формы могут быть частично не определены, с None в неизвестных измерениях. Когда задана сигнатура входных данных, tf.contrib.eager.defun будет создавать только один граф для декорированной функции Python.

import tensorflow as tf

# The first `TensorSpec` below describes the shape and dtype of `words`,
# and the second describes the shape and dtype of `another_tensor`. Note that
# the last dimension of the `words` `TensorSpec` is left unspecified.
@tf.contrib.eager.defun(input_signature=[
  tf.contrib.eager.TensorSpec(shape=[50, 300, None], dtype=tf.float32),
  tf.contrib.eager.TensorSpec(shape=[300, 100], dtype=tf.float32)
])
def my_sequence_model(words, another_tensor):
  ...

# Note how the third dimension of the first input can vary freely.
words = tf.random.uniform(([50, 300, 10])
second_input = tf.random.uniform([300, 100])
my_sequence_model(words, second_input)

words = tf.random.uniform(([50, 300, 20])
my_sequence_model(words, second_input)

# Passing an input with an incompatible shape will raise an error.
words = tf.random.uniform(([50, 100, 20])
my_sequence_model(words, second_input)  # <---- This will raise an error.

Функции Python, которые компилируются с input_signature, должны принимать только тензоры в качестве аргументов и не должны принимать именованные ключевые аргументы (**kwargs).

Отслеживание

Обратите внимание, что поскольку F регистрирует только операции TensorFlow, весь остальной код Python, который f выполняет, будет формировать только создание графов, которые F выполняет: код Python не будет выполняться при выполнении самих графов, хотя он будет выполняться каждый раз, когда функция Python отслеживается (и данная функция Python может быть отслежена несколько раз, один раз для каждой сигнатуры входных данных, с которой она вызывается). Например, тогда как функция Python

import tensorflow as tf
import numpy as np

tf.compat.v1.enable_eager_execution()

def add_noise():
  return tf.eye(5) + np.random.randn(5, 5)

будет возвращать разный результат каждый раз при вызове, скомпилированная функция compiled = tf.contrib.eager.defun(add_noise) будет возвращать одно и то же значение каждый раз при вызове, так как конкретное случайное смещение, сгенерированное NumPy, будет вставлено в граф как постоянная TensorFlow. Решение состоит в замене вызова np.random.randn на tf.random.normal((5, 5)).

Побочные эффекты Python

Следствие предыдущей дискуссии об отслеживании заключается в следующем: если функция Python f имеет побочные эффекты Python, то выполнение f несколько раз не обязательно будет семантически эквивалентно выполнению F = tf.contrib.eager.defun(f) несколько раз; это различие обусловлено тем, что defun захватывает только подграф операций TensorFlow, который создаётся, когда f вызывается в контексте построения графа.

Поток управления Python

Структура многих вычислений машинного обучения зависит от того, происходит обучение или валидация, и часто специализированная логика встраивается в блоки if training:. Картируя каждую сигнатуру входных данных на уникальный граф, defun позволяет пользователям прозрачно компилировать такой код, как демонстрирует следующий фрагмент кода:

import tensorflow as tf

tf.compat.v1.enable_eager_execution()

@tf.contrib.eager.defun
def lossy_matmul(W, x, training=True):
  outputs = tf.matmul(W, x)
  if training:
    outputs = tf.nn.dropout(outputs, keep_probability=0.2)
  return outputs

W = tf.random.normal((3, 5))
x = tf.random.normal((5, 1))

# Executes a graph that applies dropout.
lossy_outputs = lossy_matmul(W, x, training=True)

# Executes a graph that does not apply dropout.
exact_outputs = lossy_matmul(W, x, training=False)

Поток управления TensorFlow

Когда autograph True, разрешен также поток управления, зависящий от данных. Управляющие операторы, зависящие от Tensor значений, помещаются в соответствующие операции TensorFlow. Например, следующий код будет работать как ожидается:

@tf.contrib.eager.defun
def dynamic_rnn_loop(cell, seq):
  state, output = cell.zero_state()
  for input in seq:
    state, output = cell(input, state)
  return output

Для получения дополнительной информации см. tf.autograph.

Переменные

Операции TensorFlow, связанные с созданием и инициализацией переменных, автоматически выносятся за пределы графов, генерируемых defun. На практике это означает, что создание и инициализация переменных происходит только в первый раз, когда F вызывается, и что переменные повторно используются при каждом последующем вызове. Многие API TensorFlow, такие как объекты tf.keras.layers.Layer, создают переменные в первый раз при вызове и повторно используют их впоследствии. Автоматическое поднятие переменных позволяет компилировать эти API без дополнительных усилий, за счёт введения расхождения между семантикой выполнения функций Python и соответствующих скомпилированных функций. Например:

import tensorflow as tf

tf.compat.v1.enable_eager_execution()

def fn():
  x = tf.Variable(0.0)
  x.assign_add(1.0)
  return x.read_value()

# `fn` is a Python function, so x is created, initialized, and destroyed upon
# every invocation
assert fn().numpy() == fn().numpy() == 1.0

compiled = tf.contrib.eager.defun(fn)

# Compiling `fn` with `defun` hoists all variables outside of the generated
# graph, so initialization happens exactly once.
assert compiled().numpy() == 1.0
assert compiled().numpy() == 2.0

Наконец, поскольку каждая сигнатура входных данных связана с уникальным графом, если ваша функция Python создаёт tf.Variable объекты, каждый созданный для этой функции Python граф будет ссылаться на уникальный набор переменных. Чтобы обойти эту проблему, мы рекомендуем не компилировать функции Python, создающие tf.Variable объекты. Вместо этого функции Python должны либо лексически ссылаться на tf.Variable объекты, либо принимать их в качестве аргументов, предпочтительно в объектно-ориентированном контейнере. Если вам необходимо создавать переменные внутри вашей функции Python, и вы хотите, чтобы каждый сгенерированный для неё граф ссылался на один и тот же набор переменных, добавьте в вашу функцию Python логику, которая гарантирует, что переменные создаются только в первый раз при её вызове и повторно используются для каждого последующего вызова; обратите внимание, что именно это делают объекты tf.keras.layers.Layer, поэтому мы рекомендуем использовать их для представления вычислений, связанных с переменными, по возможности.