tf.function

Создаёт вызываемый граф TensorFlow из Python-функции.

tf.function(
    func=None, input_signature=None, autograph=True,
    experimental_autograph_options=None, experimental_relax_shapes=False,
    experimental_compile=None
)

function создаёт вызываемый объект, который выполняет граф TensorFlow (tf.Graph), созданный путём отслеживания операций TensorFlow в func. Это позволяет исполняющей системе TensorFlow применять оптимизации и использовать параллелизм в вычислениях, определённых func.

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

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

g = tf.function(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()

# Tensors and tf.Variables used by the Python function are captured in the
# graph.
@tf.function
def h():
  return f(x, y)

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

# Data-dependent control flow is also captured in the graph. Supported
# control flow statements include `if`, `for`, `while`, `break`, `continue`,
# `return`.
@tf.function
def g(x):
  if tf.reduce_sum(x) > 0:
    return x * x
  else:
    return -x // 2

# print and TensorFlow side effects are supported, but exercise caution when
# using Python side effects like mutating objects, saving to files, etc.
l = []

@tf.function
def g(x):
  for i in x:
    print(i)                              # Works
    tf.compat.v1.assign(v, i)                       # Works
    tf.compat.v1.py_func(lambda i: l.append(i))(i)  # Works
    l.append(i)                           # Caution! Doesn't work.

Обратите внимание, что в отличие от других операций TensorFlow, мы не конвертируем числовые входные данные Python в тензоры. Кроме того, для каждого уникального числового значения Python генерируется новый граф. Например, вызовы g(2) и g(3) сгенерируют два новых графа (в то время как один граф сгенерируется при вызове g(tf.constant(2)) и g(tf.constant(3))). Поэтому числовые входные данные Python должны быть ограничены аргументами, которые будут иметь мало уникальных значений, такими как гиперпараметры, например, количество слоёв в нейронной сети. Это позволяет TensorFlow оптимизировать каждый вариант нейронной сети.

Ссылка на tf.Variable

Python-функция func может ссылаться на объекты состояния (например, tf.Variable). Они фиксируются в качестве неявных входов для вызываемого объекта, возвращаемого function. Например:

c = tf.Variable(0)

@tf.function
def f(x):
  c.assign_add(1)
  return x + tf.compat.v1.to_float(c)

assert int(c) == 0
assert f(1.0) == 2.0
assert int(c) == 1
assert f(1.0) == 3.0
assert int(c) == 2

function может быть применено к методам объекта. Например:

class Dense(object):
  def __init__(self):
    self.W = tf.Variable(tf.compat.v1.glorot_uniform_initializer()((10, 10)))
    self.b = tf.Variable(tf.zeros(10))

  @tf.function
  def compute(self, x):
    return tf.matmul(x, self.W) + self.b

d1 = Dense()
d2 = Dense()
x = tf.random.uniform((10, 10))
# d1 and d2 are using distinct variables
assert not (d1.compute(x).numpy() == d2.compute(x).numpy()).all()

Использование с tf.keras

Методы call подкласса tf.keras.Model могут быть декорированы function для применения оптимизаций выполнения графа к ним. Например:

class MyModel(tf.keras.Model):
  def __init__(self, keep_probability=0.2):
    super(MyModel, self).__init__()
    self.dense1 = tf.keras.layers.Dense(4)
    self.dense2 = tf.keras.layers.Dense(5)
    self.keep_probability = keep_probability

  @tf.function
  def call(self, inputs, training=True):
    y = self.dense2(self.dense1(inputs))
    if training:
      return tf.nn.dropout(y, self.keep_probability)
    else:
      return y

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

Подписи входных данных

function создаёт отдельный граф для каждой уникальной комбинации форм и типов входных данных. Например, следующий фрагмент кода приведёт к построению трёх отдельных графов, поскольку у каждого входа разная форма.

@tf.function
def f(x): return tf.add(x, 1.)

scalar = tf.constant(1.0)
vector = tf.constant([1.0, 1.0])
matrix = tf.constant([[3.0]])

f(scalar)
f(vector)
f(matrix)

«Подпись входных данных» может быть необязательно предоставлена для function для управления отслеживаемыми графами. Подпись входных данных определяет форму и тип каждого аргумента Tensor функции, используя объект tf.TensorSpec. Например, следующий фрагмент кода гарантирует создание одного графа, где вход Tensor должен быть тензором с плавающей точкой без ограничений по форме.

@tf.function(input_signature=[tf.TensorSpec(shape=None, dtype=tf.float32)])
def f(x): return tf.add(x, 1.)

Когда указана подпись, вызываемый объект преобразует входные данные в указанные TensorSpecs.

Отслеживание и постановка

Когда autograph равно True, весь Python-управляемый поток, зависящий от значений Tensor, ставится в граф TensorFlow. Когда autograph равно False, функция отслеживается, и управляемому потоку не разрешается зависеть от данных.

Обратите внимание, что function ставит в очередь только операции TensorFlow, весь Python-код, который func выполняет и не зависит от данных, будет определять построение графа. Например, рассмотрите следующее:

import numpy as np

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

traced = tf.function(add_noise)

add_noise() будет возвращать разные результаты каждый раз при вызове. Однако, traced() будет возвращать одно и то же значение каждый раз при вызове, поскольку конкретное случайное значение, сгенерированное вызовом np.random.randn, будет вставлено в отслеженный/поставленный граф TensorFlow в качестве константы. В этом конкретном примере замена np.random.randn(5, 5) на tf.random.normal((5, 5)) приведёт к такому же поведению для add_noise() и traced().

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

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

То же самое относится к коду с побочными эффектами на стороне Python, используемому внутри управляемого потока, например, в цикле. Если ваш код использует побочные эффекты, которые не предназначены для управления построением графа, оберните их в tf.compat.v1.py_func.

Переотслеживание

Один объект tf.function может потребовать сопоставления с несколькими вычислительными графами под капотом. Это должно быть видно только как производительность (отслеживание графов имеет ненулевую вычислительную и ресурсную стоимость), но не должно влиять на правильность программы. Отслеженная функция должна возвращать тот же результат, что и при выполнении её в режиме «встречного», при условии отсутствия побочных эффектов на стороне Python.

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

Помимо этого, TensorFlow оставляет за собой право переотслеживать функции столько раз, сколько необходимо, для обеспечения того, чтобы отслеженные функции вели себя так же, как и при выполнении их в режиме «встречного», и для обеспечения наилучшей общей производительности. Например, поведение, определяющее, сколько прослеживаний выполнит TensorFlow при многократном вызове функции с различными числовыми аргументами в качестве Python-аргументов, не определено для возможности будущих оптимизаций.

Для управления поведением отслеживания используйте следующие инструменты:

• разные tf.function объекты гарантированно не будут использовать общие следы; и
• указание подписи или использование конкретных объектов функций, возвращаемых get_concrete_function(), гарантирует, что будет создан только один граф функции.