Se usó la API de Cloud Translation para traducir esta página.
Switch to English

tf.function

Versión TensorFlow 1 Ver código fuente en GitHub

Compila una función en un gráfico invocable TensorFlow.

Usado en los cuadernos

Usado en la guía Usado en los tutoriales

tf.function construye un invocable que ejecuta un gráfico TensorFlow ( tf.Graph ) creado compilando las operaciones de TensorFlow en func , ejecutando efectivamente func como un gráfico TensorFlow.

Ejemplo de uso:

@tf.function
def f(x, y):
  return x ** 2 + y
x = tf.constant([2, 3])
y = tf.constant([3, -2])
f(x, y)
<tf.Tensor: ... numpy=array([7, 7], ...)>

Caracteristicas

func puede usar un flujo de control dependiente de los datos, incluyendo if , for , while break , continue y return declaraciones:

@tf.function
def f(x):
  if tf.reduce_sum(x) > 0:
    return x * x
  else:
    return -x // 2
f(tf.constant(-2))
<tf.Tensor: ... numpy=1>

El cierre de func puede incluir tf.Tensor y tf.Variable Objetos tf.Variable :

@tf.function
def f():
  return x ** 2 + y
x = tf.constant([-2, -3])
y = tf.Variable([3, -2])
f()
<tf.Tensor: ... numpy=array([7, 7], ...)>

func también puede usar operaciones con efectos secundarios, como tf.print , tf.Variable y otros:

v = tf.Variable(1)
@tf.function
def f(x):
  for i in tf.range(x):
    v.assign_add(i)
f(3)
v
<tf.Variable ... numpy=4>
l = []
@tf.function
def f(x):
  for i in x:
    l.append(i + 1)    # Caution! Will only happen once when tracing
f(tf.constant([1, 2, 3]))
l
[<tf.Tensor ...>]

En su lugar, use colecciones TensorFlow como tf.TensorArray :

@tf.function
def f(x):
  ta = tf.TensorArray(dtype=tf.int32, size=0, dynamic_size=True)
  for i in range(len(x)):
    ta = ta.write(i, x[i] + 1)
  return ta.stack()
f(tf.constant([1, 2, 3]))
<tf.Tensor: ..., numpy=array([2, 3, 4], ...)>

tf.function es polimórfica

Internamente, tf.function puede construir más de un gráfico, para admitir argumentos con diferentes tipos de datos o formas, ya que TensorFlow puede construir gráficos más eficientes que se especializan en formas y tipos. tf.function también trata cualquier valor de Python puro como objetos opacos, y construye un gráfico separado para cada conjunto de argumentos de Python que encuentra.

Para obtener un gráfico individual, use el método get_concrete_function del invocable creado por tf.function . Se puede tf.Graph con los mismos argumentos que func y devuelve un objeto tf.Graph especial:

@tf.function
def f(x):
  return x + 1
isinstance(f.get_concrete_function(1).graph, tf.Graph)
True
@tf.function
def f(x):
  return tf.abs(x)
f1 = f.get_concrete_function(1)
f2 = f.get_concrete_function(2)  # Slow - builds new graph
f1 is f2
False
f1 = f.get_concrete_function(tf.constant(1))
f2 = f.get_concrete_function(tf.constant(2))  # Fast - reuses f1
f1 is f2
True

Los argumentos numéricos de Python solo deben usarse cuando toman pocos valores distintos, como hiperparámetros como el número de capas en una red neuronal.

Firmas de entrada

Para los argumentos de Tensor, la función tf.function una instancia de un gráfico separado para cada conjunto único de formas de entrada y tipos de datos. El siguiente ejemplo crea dos gráficos separados, cada uno especializado en una forma diferente:

@tf.function
def f(x):
  return x + 1
vector = tf.constant([1.0, 1.0])
matrix = tf.constant([[3.0]])
f.get_concrete_function(vector) is f.get_concrete_function(matrix)
False

Se puede proporcionar opcionalmente una "firma de entrada" a tf.function para controlar los gráficos trazados. La firma de entrada especifica la forma y el tipo de cada argumento Tensor para la función utilizando un objeto tf.TensorSpec . Se pueden usar formas más generales. Esto es útil para evitar crear múltiples gráficos cuando los tensores tienen formas dinámicas. También restringe la forma y el tipo de datos de los tensores que se pueden usar:

@tf.function(
    input_signature=[tf.TensorSpec(shape=None, dtype=tf.float32