tf_numpy.ipynb

Kernel: Python 3

Copyright 2020 The TensorFlow Authors.

#@title Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

API NumPy en TensorFlow

Ver en TensorFlow.org

Ejecutar en Google Colab

Ver código fuente en GitHub

Descargar el bloc de notas

Descripción general

TensorFlow implementa un subconjunto de la API NumPy, disponible como tf.experimental.numpy. Esto permite ejecutar código NumPy, acelerado por TensorFlow, y al mismo tiempo permite el acceso a todas las API de TensorFlow.

Preparación

import matplotlib.pyplot as plt
import numpy as np
import tensorflow as tf
import tensorflow.experimental.numpy as tnp
import timeit

print("Using TensorFlow version %s" % tf.__version__)

Habilitar el comportamiento NumPy

Para utilizar tnp como NumPy, habilite el comportamiento de NumPy para TensorFlow:

tnp.experimental_enable_numpy_behavior()

Esta llamada habilita la promoción de tipos en TensorFlow y también cambia la inferencia de tipos, al convertir literales en tensores, para seguir más estrictamente el estándar NumPy.

Nota: Esta llamada cambiará el comportamiento de todo TensorFlow, no solo el módulo tf.experimental.numpy.

Arreglo NumPy ND de TensorFlow

Una instancia de tf.experimental.numpy.ndarray, que se llama arreglo ND, representa un arreglo denso multidimensional de un dtype determinado en un determinado dispositivo. Es un alias de tf.Tensor. Consulte la clase de arreglo ND para conocer métodos útiles como ndarray.T, ndarray.reshape, ndarray.ravel y otros.

Primero cree un objeto de arreglo ND y luego invoque diferentes métodos.

# Create an ND array and check out different attributes.
ones = tnp.ones([5, 3], dtype=tnp.float32)
print("Created ND array with shape = %s, rank = %s, "
      "dtype = %s on device = %s\n" % (
          ones.shape, ones.ndim, ones.dtype, ones.device))

# `ndarray` is just an alias to `tf.Tensor`.
print("Is `ones` an instance of tf.Tensor: %s\n" % isinstance(ones, tf.Tensor))

# Try commonly used member functions.
print("ndarray.T has shape %s" % str(ones.T.shape))
print("narray.reshape(-1) has shape %s" % ones.reshape(-1).shape)

Promoción de tipos

Hay 4 opciones de promoción de tipos en TensorFlow.

• Por defecto, TensorFlow arroja errores en vez de promover tipos para operaciones de tipo mezclado.

• Al ejecutar tf.numpy.experimental_enable_numpy_behavior(), TensorFlow cambia y usa reglas de promoción de tipos NumPy.

• Después de TensorFlow 2.15, hay dos opciones nuevas (consulte Promoción de tipo TF NumPy para obtener más detalles):

  • tf.numpy.experimental_enable_numpy_behavior() usa reglas de promoción de tipo Jax

  • tf.numpy.experimental_enable_numpy_behavior(dtype_conversion_mode="safe") usa reglas de promoción de tipo Jax, pero no permite ciertas promociones que no son seguras.

Promoción de tipo NumPy

Las API de TensorFlow NumPy tienen una semántica bien definida para convertir literales a arreglos ND, así como para realizar promoción de tipos en entradas de arreglos ND. Consultenp.result_type para obtener más detalles.

Las API de TensorFlow no hacen cambios en las entradas tf.Tensor y no realizan promoción de tipos en ellas, mientras que las API de TensorFlow NumPy promueven todas las entradas de acuerdo con las reglas de promoción de tipos de NumPy. En el siguiente ejemplo, realizará una promoción de tipo. Primero, ejecute la suma en entradas de arreglo ND de diferentes tipos y observe los tipos de salida. Las API de TensorFlow no permitirían ninguna promoción de este tipo.

print("Type promotion for operations")
values = [tnp.asarray(1, dtype=d) for d in
          (tnp.int32, tnp.int64, tnp.float32, tnp.float64)]
for i, v1 in enumerate(values):
  for v2 in values[i + 1:]:
    print("%s + %s => %s" %
          (v1.dtype.name, v2.dtype.name, (v1 + v2).dtype.name))

Por último, convierta literales a un arreglo ND con ndarray.asarray y observe el tipo que resulta.

print("Type inference during array creation")
print("tnp.asarray(1).dtype == tnp.%s" % tnp.asarray(1).dtype.name)
print("tnp.asarray(1.).dtype == tnp.%s\n" % tnp.asarray(1.).dtype.name)

Al convertir literales a arreglos ND, NumPy prefiere tipos anchos como tnp.int64 y tnp.float64. Por el contrario, tf.convert_to_tensor prefiere los tipos tf.int32 y tf.float32 para convertir constantes a tf.Tensor. Las API de TensorFlow NumPy se adhieren al comportamiento de NumPy para números enteros. En cuanto a los flotadores, el argumento prefer_float32 de experimental_enable_numpy_behavior le permite controlar si prefiere tf.float32 sobre tf.float64 (el valor predeterminado es False). Por ejemplo:

tnp.experimental_enable_numpy_behavior(prefer_float32=True)
print("When prefer_float32 is True:")
print("tnp.asarray(1.).dtype == tnp.%s" % tnp.asarray(1.).dtype.name)
print("tnp.add(1., 2.).dtype == tnp.%s" % tnp.add(1., 2.).dtype.name)

tnp.experimental_enable_numpy_behavior(prefer_float32=False)
print("When prefer_float32 is False:")
print("tnp.asarray(1.).dtype == tnp.%s" % tnp.asarray(1.).dtype.name)
print("tnp.add(1., 2.).dtype == tnp.%s" % tnp.add(1., 2.).dtype.name)

Transmisión

Al igual que TensorFlow, NumPy define una semántica rica para valores de "difusión". Puede consultar la guía de transmisión de NumPy para obtener más información y compararla con la semántica de transmisión de TensorFlow.

x = tnp.ones([2, 3])
y = tnp.ones([3])
z = tnp.ones([1, 2, 1])
print("Broadcasting shapes %s, %s and %s gives shape %s" % (
    x.shape, y.shape, z.shape, (x + y + z).shape))

Indexación

NumPy define reglas de indexación muy sofisticadas. Consulte la guía de indexación de NumPy. Observe el uso de arreglos ND como índices a continuación.

x = tnp.arange(24).reshape(2, 3, 4)

print("Basic indexing")
print(x[1, tnp.newaxis, 1:3, ...], "\n")

print("Boolean indexing")
print(x[:, (True, False, True)], "\n")

print("Advanced indexing")
print(x[1, (0, 0, 1), tnp.asarray([0, 1, 1])])

# Mutation is currently not supported
try:
  tnp.arange(6)[1] = -1
except TypeError:
  print("Currently, TensorFlow NumPy does not support mutation.")

Modelo de ejemplo

A continuación, puede ver cómo crear un modelo y ejecutar inferencias sobre él. Este modelo simple aplica una capa relu seguida de una proyección lineal. En las secciones posteriores se mostrará cómo calcular gradientes para este modelo con GradientTape de TensorFlow.

class Model(object):
  """Model with a dense and a linear layer."""

  def __init__(self):
    self.weights = None

  def predict(self, inputs):
    if self.weights is None:
      size = inputs.shape[1]
      # Note that type `tnp.float32` is used for performance.
      stddev = tnp.sqrt(size).astype(tnp.float32)
      w1 = tnp.random.randn(size, 64).astype(tnp.float32) / stddev
      bias = tnp.random.randn(64).astype(tnp.float32)
      w2 = tnp.random.randn(64, 2).astype(tnp.float32) / 8
      self.weights = (w1, bias, w2)
    else:
      w1, bias, w2 = self.weights
    y = tnp.matmul(inputs, w1) + bias
    y = tnp.maximum(y, 0)  # Relu
    return tnp.matmul(y, w2)  # Linear projection

model = Model()
# Create input data and compute predictions.
print(model.predict(tnp.ones([2, 32], dtype=tnp.float32)))

TensorFlow NumPy y NumPy

TensorFlow NumPy implementa un subconjunto de la especificación NumPy completa. Si bien se agregarán más símbolos con el tiempo, hay funciones sistemáticas que no se admitirán en un futuro próximo. Entre ellos se incluyen la compatibilidad con NumPy C API, la integración de Swig, el orden de almacenamiento de Fortran, las vistas y stride_tricks y algunos dtype (como np.recarray y np.object). Para obtener más detalles, consulte la documentación de la API TensorFlow NumPy.

Interoperabilidad numérica

Los arreglos TensorFlow ND pueden interoperar con funciones NumPy. Estos objetos implementan la interfaz __array__. NumPy usa esta interfaz para convertir argumentos de función a valores np.ndarray antes de procesarlos.

De manera similar, las funciones de TensorFlow NumPy pueden aceptar entradas de diferentes tipos, incluido np.ndarray. Estas entradas se convierten en un arreglo ND al llamar a ndarray.asarray.

La conversión del arreglo ND hacia y desde np.ndarray puede generar copias de datos reales. Consulte la sección sobre copias en búfer para obtener más detalles.

# ND array passed into NumPy function.
np_sum = np.sum(tnp.ones([2, 3]))
print("sum = %s. Class: %s" % (float(np_sum), np_sum.__class__))

# `np.ndarray` passed into TensorFlow NumPy function.
tnp_sum = tnp.sum(np.ones([2, 3]))
print("sum = %s. Class: %s" % (float(tnp_sum), tnp_sum.__class__))

# It is easy to plot ND arrays, given the __array__ interface.
labels = 15 + 2 * tnp.random.randn(1, 1000)
_ = plt.hist(labels)

Copias de búfer

Mezclar TensorFlow NumPy con código NumPy puede desencadenar copias de datos. Esto se debe a que TensorFlow NumPy tiene requisitos de alineación de memoria más estrictos que los de NumPy.

Cuando se pasa un np.ndarray a TensorFlow NumPy, se comprobarán los requisitos de alineación y se activará una copia si es necesario. Al pasar un búfer de CPU de arreglo ND a NumPy, por lo general, el búfer cumplirá los requisitos de alineación y NumPy no necesitará crear una copia.

Los arreglos ND pueden hacer referencia a búferes ubicados en dispositivos distintos de la memoria de la CPU local. En tales casos, llamar a una función NumPy activará copias en la red o en los dispositivo según sea necesario.

Por esto, la mezcla con llamadas a la API de NumPy generalmente deben realizarse con precaución y el usuario debe tener cuidado con los gastos generales de copia de datos. Generalmente, intercalar llamadas de TensorFlow NumPy con llamadas de TensorFlow es seguro y evita la copia de datos. Consulte la sección sobre interoperabilidad de TensorFlow para obtener más detalles.

Prioridad del operador

TensorFlow NumPy define una __array_priority__ mayor que la de NumPy. Esto significa que para los operadores que involucran tanto un arreglo ND como np.ndarray, el primero tendrá prioridad, es decir, la entrada de np.ndarray se convertirá en un arreglo ND y se llamará a la implementación TensorFlow NumPy del operador.

x = tnp.ones([2]) + np.ones([2])
print("x = %s\nclass = %s" % (x, x.__class__))

TF NumPy y TensorFlow

TensorFlow NumPy está construido sobre TensorFlow y, por lo tanto, interopera perfectamente con TensorFlow.

tf.Tensor y arreglo ND

El arreglo ND es un alias de tf.Tensor, por lo que obviamente se pueden mezclar sin activar copias de datos reales.

x = tf.constant([1, 2])
print(x)

# `asarray` and `convert_to_tensor` here are no-ops.
tnp_x = tnp.asarray(x)
print(tnp_x)
print(tf.convert_to_tensor(tnp_x))

# Note that tf.Tensor.numpy() will continue to return `np.ndarray`.
print(x.numpy(), x.numpy().__class__)

Interoperabilidad de TensorFlow

Se puede pasar un arreglo ND a las API de TensorFlow, ya que el arreglo ND es solo un alias de tf.Tensor. Como se mencionó anteriormente, dicha interoperación no realiza copias de datos, ni siquiera de datos ubicados en aceleradores ni dispositivos remotos.

Por el contrario, los objetos tf.Tensor se pueden pasar a las API tf.experimental.numpy sin realizar copias de datos.

# ND array passed into TensorFlow function.
tf_sum = tf.reduce_sum(tnp.ones([2, 3], tnp.float32))
print("Output = %s" % tf_sum)

# `tf.Tensor` passed into TensorFlow NumPy function.
tnp_sum = tnp.sum(tf.ones([2, 3]))
print("Output = %s" % tnp_sum)

Gradientes y jacobianos: tf.GradientTape

GradientTape de TensorFlow se puede usar para la retropropagación a través del código de TensorFlow y TensorFlow NumPy.

Use el modelo creado en la sección Modelo de ejemplo y calcule gradientes y jacobianos.

def create_batch(batch_size=32):
  """Creates a batch of input and labels."""
  return (tnp.random.randn(batch_size, 32).astype(tnp.float32),
          tnp.random.randn(batch_size, 2).astype(tnp.float32))

def compute_gradients(model, inputs, labels):
  """Computes gradients of squared loss between model prediction and labels."""
  with tf.GradientTape() as tape:
    assert model.weights is not None
    # Note that `model.weights` need to be explicitly watched since they
    # are not tf.Variables.
    tape.watch(model.weights)
    # Compute prediction and loss
    prediction = model.predict(inputs)
    loss = tnp.sum(tnp.square(prediction - labels))
  # This call computes the gradient through the computation above.
  return tape.gradient(loss, model.weights)

inputs, labels = create_batch()
gradients = compute_gradients(model, inputs, labels)

# Inspect the shapes of returned gradients to verify they match the
# parameter shapes.
print("Parameter shapes:", [w.shape for w in model.weights])
print("Gradient shapes:", [g.shape for g in gradients])
# Verify that gradients are of type ND array.
assert isinstance(gradients[0], tnp.ndarray)

# Computes a batch of jacobians. Each row is the jacobian of an element in the
# batch of outputs w.r.t. the corresponding input batch element.
def prediction_batch_jacobian(inputs):
  with tf.GradientTape() as tape:
    tape.watch(inputs)
    prediction = model.predict(inputs)
  return prediction, tape.batch_jacobian(prediction, inputs)

inp_batch = tnp.ones([16, 32], tnp.float32)
output, batch_jacobian = prediction_batch_jacobian(inp_batch)
# Note how the batch jacobian shape relates to the input and output shapes.
print("Output shape: %s, input shape: %s" % (output.shape, inp_batch.shape))
print("Batch jacobian shape:", batch_jacobian.shape)

Compilación de trazado: tf.function

tf.function de TensorFlow funciona mediante la "compilación de trazado" del código y luego optimiza estos trazados para un rendimiento mucho más rápido. Consulte la Introducción a gráficos y funciones.

tf.function también se puede usar para optimizar el código de TensorFlow NumPy. Aquí hay un ejemplo simple para demostrar las aceleraciones. Tenga en cuenta que el cuerpo del código tf.function incluye llamadas a las API de TensorFlow NumPy.

inputs, labels = create_batch(512)
print("Eager performance")
compute_gradients(model, inputs, labels)
print(timeit.timeit(lambda: compute_gradients(model, inputs, labels),
                    number=10) * 100, "ms")

print("\ntf.function compiled performance")
compiled_compute_gradients = tf.function(compute_gradients)
compiled_compute_gradients(model, inputs, labels)  # warmup
print(timeit.timeit(lambda: compiled_compute_gradients(model, inputs, labels),
                    number=10) * 100, "ms")

Vectorización: tf.vectorized_map

TensorFlow tiene soporte incorporado para vectorizar bucles paralelos, lo que permite aceleraciones de uno a dos órdenes de magnitud. Se puede acceder a estas aceleraciones a través de la API tf.vectorized_map y también se aplican al código TensorFlow NumPy.

A veces es útil calcular el gradiente de cada salida en un lote con respecto al elemento del lote de entrada correspondiente. Dicho cálculo se puede realizar de manera eficiente con tf.vectorized_map como se muestra a continuación.

@tf.function
def vectorized_per_example_gradients(inputs, labels):
  def single_example_gradient(arg):
    inp, label = arg
    return compute_gradients(model,
                             tnp.expand_dims(inp, 0),
                             tnp.expand_dims(label, 0))
  # Note that a call to `tf.vectorized_map` semantically maps
  # `single_example_gradient` over each row of `inputs` and `labels`.
  # The interface is similar to `tf.map_fn`.
  # The underlying machinery vectorizes away this map loop which gives
  # nice speedups.
  return tf.vectorized_map(single_example_gradient, (inputs, labels))

batch_size = 128
inputs, labels = create_batch(batch_size)

per_example_gradients = vectorized_per_example_gradients(inputs, labels)
for w, p in zip(model.weights, per_example_gradients):
  print("Weight shape: %s, batch size: %s, per example gradient shape: %s " % (
      w.shape, batch_size, p.shape))

# Benchmark the vectorized computation above and compare with
# unvectorized sequential computation using `tf.map_fn`.
@tf.function
def unvectorized_per_example_gradients(inputs, labels):
  def single_example_gradient(arg):
    inp, label = arg
    return compute_gradients(model,
                             tnp.expand_dims(inp, 0),
                             tnp.expand_dims(label, 0))

  return tf.map_fn(single_example_gradient, (inputs, labels),
                   fn_output_signature=(tf.float32, tf.float32, tf.float32))

print("Running vectorized computation")
print(timeit.timeit(lambda: vectorized_per_example_gradients(inputs, labels),
                    number=10) * 100, "ms")

print("\nRunning unvectorized computation")
per_example_gradients = unvectorized_per_example_gradients(inputs, labels)
print(timeit.timeit(lambda: unvectorized_per_example_gradients(inputs, labels),
                    number=10) * 100, "ms")

Colocación del dispositivo

TensorFlow NumPy puede realizar operaciones en CPU, GPU, TPU y dispositivos remotos. Usa mecanismos estándar de TensorFlow para la colocación del dispositivo. A continuación, un ejemplo simple muestra cómo enumerar todos los dispositivos y luego realizar algunos cálculos en un dispositivo en particular.

TensorFlow también tiene API para replicar cálculos entre dispositivos y realizar reducciones colectivas que no trataremos aquí.

Listar dispositivos

tf.config.list_logical_devices y tf.config.list_physical_devices se pueden usar para encontrar qué dispositivos usar.

print("All logical devices:", tf.config.list_logical_devices())
print("All physical devices:", tf.config.list_physical_devices())

# Try to get the GPU device. If unavailable, fallback to CPU.
try:
  device = tf.config.list_logical_devices(device_type="GPU")[0]
except IndexError:
  device = "/device:CPU:0"

Operaciones de apartado: tf.device

Se pueden apartar las operaciones en un dispositivo al llamarlo en un alcance tf.device.

print("Using device: %s" % str(device))
# Run operations in the `tf.device` scope.
# If a GPU is available, these operations execute on the GPU and outputs are
# placed on the GPU memory.
with tf.device(device):
  prediction = model.predict(create_batch(5)[0])

print("prediction is placed on %s" % prediction.device)

Copiar arreglos ND entre dispositivos: tnp.copy

Una llamada a tnp.copy, apartada en un determinado alcance de dispositivo, copiará los datos a ese dispositivo, a menos que los datos ya estén en ese dispositivo.

with tf.device("/device:CPU:0"):
  prediction_cpu = tnp.copy(prediction)
print(prediction.device)
print(prediction_cpu.device)

Comparaciones de rendimiento

TensorFlow NumPy usa núcleos de TensorFlow altamente optimizados que se pueden distribuir en CPU, GPU y TPU. TensorFlow también realiza muchas optimizaciones del compilador, como la fusión de operaciones, que se traducen en mejoras de rendimiento y memoria. Consulte Optimización de gráficos de TensorFlow con Grappler para obtener más información.

Sin embargo, TensorFlow tiene mayores gastos generales para las operaciones de envío en comparación con NumPy. Para cargas de trabajo compuestas por operaciones pequeñas (menos de aproximadamente 10 microsegundos), estos gastos generales pueden dominar el tiempo de ejecución y NumPy podría proporcionar un mejor rendimiento. Para otros casos, TensorFlow generalmente debería proporcionar un mejor rendimiento.

Ejecute el siguiente punto de referencia para comparar el rendimiento de NumPy y TensorFlow NumPy para diferentes tamaños de entrada.

def benchmark(f, inputs, number=30, force_gpu_sync=False):
  """Utility to benchmark `f` on each value in `inputs`."""
  times = []
  for inp in inputs:
    def _g():
      if force_gpu_sync:
        one = tnp.asarray(1)
      f(inp)
      if force_gpu_sync:
        with tf.device("CPU:0"):
          tnp.copy(one)  # Force a sync for GPU case

    _g()  # warmup
    t = timeit.timeit(_g, number=number)
    times.append(t * 1000. / number)
  return times


def plot(np_times, tnp_times, compiled_tnp_times, has_gpu, tnp_times_gpu):
  """Plot the different runtimes."""
  plt.xlabel("size")
  plt.ylabel("time (ms)")
  plt.title("Sigmoid benchmark: TF NumPy vs NumPy")
  plt.plot(sizes, np_times, label="NumPy")
  plt.plot(sizes, tnp_times, label="TF NumPy (CPU)")
  plt.plot(sizes, compiled_tnp_times, label="Compiled TF NumPy (CPU)")
  if has_gpu:
    plt.plot(sizes, tnp_times_gpu, label="TF NumPy (GPU)")
  plt.legend()

# Define a simple implementation of `sigmoid`, and benchmark it using
# NumPy and TensorFlow NumPy for different input sizes.

def np_sigmoid(y):
  return 1. / (1. + np.exp(-y))

def tnp_sigmoid(y):
  return 1. / (1. + tnp.exp(-y))

@tf.function
def compiled_tnp_sigmoid(y):
  return tnp_sigmoid(y)

sizes = (2 ** 0, 2 ** 5, 2 ** 10, 2 ** 15, 2 ** 20)
np_inputs = [np.random.randn(size).astype(np.float32) for size in sizes]
np_times = benchmark(np_sigmoid, np_inputs)

with tf.device("/device:CPU:0"):
  tnp_inputs = [tnp.random.randn(size).astype(np.float32) for size in sizes]
  tnp_times = benchmark(tnp_sigmoid, tnp_inputs)
  compiled_tnp_times = benchmark(compiled_tnp_sigmoid, tnp_inputs)

has_gpu = len(tf.config.list_logical_devices("GPU"))
if has_gpu:
  with tf.device("/device:GPU:0"):
    tnp_inputs = [tnp.random.randn(size).astype(np.float32) for size in sizes]
    tnp_times_gpu = benchmark(compiled_tnp_sigmoid, tnp_inputs, 100, True)
else:
  tnp_times_gpu = None
plot(np_times, tnp_times, compiled_tnp_times, has_gpu, tnp_times_gpu)

Otras lecturas

• TensorFlow NumPy: Tutorial de clasificación de imágenes distribuidas

• TensorFlow NumPy: Keras y estrategia de distribución

• Análisis de sentimiento con Trax y TensorFlow NumPy