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

API de NumPy en TensorFlow

Ver en TensorFlow.org Ejecutar en Google Colab Ver fuente en GitHub Descargar cuaderno

Visión general

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

Preparar

pip install --quiet --upgrade tf-nightly
WARNING: You are using pip version 20.2.2; however, version 20.2.3 is available.
You should consider upgrading via the '/tmpfs/src/tf_docs_env/bin/python -m pip install --upgrade pip' command.

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__)
Using TensorFlow version 2.4.0-dev20200922

Matriz ND de TensorFlow NumPy

Una instancia de tf.experimental.numpy.ndarray , llamada ND Array , representa una matriz densa multidimensional de un tipo dtype dado colocado en un dispositivo determinado. Cada uno de estos objetos envuelve internamente un tf.Tensor . Consulte la clase de matriz ND para obtener métodos útiles como ndarray.T , ndarray.reshape , ndarray.ravel y otros.

Primero cree un objeto de matriz 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.data.device))

# Check out the internally wrapped `tf.Tensor` object.
print("The ND array wraps a tf.Tensor: %s\n" % ones.data)

# 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)
Created ND array with shape = (5, 3), rank = 2, dtype = float32 on device = /job:localhost/replica:0/task:0/device:CPU:0

The ND array wraps a tf.Tensor: tf.Tensor(
[[1. 1. 1.]
 [1. 1. 1.]
 [1. 1. 1.]
 [1. 1. 1.]
 [1. 1. 1.]], shape=(5, 3), dtype=float32)

ndarray.T has shape (3, 5)
narray.reshape(-1) has shape 15

Promoción de tipo

Las API de TensorFlow NumPy tienen una semántica bien definida para convertir literales en una matriz ND, así como para realizar la promoción de tipo en las entradas de la matriz ND. Consulte np.result_type para obtener más detalles. . Al convertir literales a una matriz 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 dejan tf.Tensor entradas tf.Tensor sin cambios y no realizan promoción de tipo en ellas.

En el siguiente ejemplo, realizará una promoción de tipo. Primero, ejecute la suma en las entradas de la matriz ND de diferentes tipos y observe los tipos de salida. Ninguna de estas promociones de tipo se permitiría en objetos tf.Tensor directos. Finalmente, convierta literales a una matriz ND usando ndarray.asarrray y observe el tipo resultante.

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, v2.dtype, (v1 + v2).dtype))

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

Type promotion for operations
int32 + int64 => int64
int32 + float32 => float64
int32 + float64 => float64
int64 + float32 => float64
int64 + float64 => float64
float32 + float64 => float64
Type inference during array creation
tnp.asarray(1).dtype == tnp.int64
tnp.asarray(1.).dtype == tnp.float64


Radiodifusión

Similar a TensorFlow, NumPy define semántica enriquecida para valores de "transmisió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))
Broadcasting shapes (2, 3), (3,) and (1, 2, 1) gives shape (1, 2, 3)

Indexación

NumPy define reglas de indexación muy sofisticadas. Consulte la guía de indexación de NumPy . Tenga en cuenta el uso de matrices 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])])
Basic indexing
ndarray<tf.Tensor(
[[[16 17 18 19]
  [20 21 22 23]]], shape=(1, 2, 4), dtype=int64)> 

Boolean indexing
ndarray<tf.Tensor(
[[[ 0  1  2  3]
  [ 8  9 10 11]]

 [[12 13 14 15]
  [20 21 22 23]]], shape=(2, 2, 4), dtype=int64)> 

Advanced indexing
ndarray<tf.Tensor([12 13 17], shape=(3,), dtype=int64)>

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

Modelo de ejemplo

A continuación, puede ver cómo crear un modelo y ejecutar inferencias en él. Este modelo simple aplica una capa relu seguida de una proyección lineal. Las secciones posteriores mostrarán 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)))
ndarray<tf.Tensor(
[[-0.31255645  0.00103381]
 [-0.31255645  0.00103381]], shape=(2, 2), dtype=float32)>

TensorFlow NumPy y NumPy

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

Interoperabilidad NumPy

Las matrices de 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, incluidos tf.Tensor y np.ndarray . Estas entradas se convierten en una matriz ND llamando a ndarray.asarray en ellas.

La conversión de la matriz ND hacia y desde np.ndarray puede desencadenar 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__))
sum = 6.0. Class: <class 'numpy.float64'>
sum = 6.0. Class: <class 'tensorflow.python.ops.numpy_ops.np_arrays.ndarray'>

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

png

Copias de búfer

La combinación de TensorFlow NumPy con código NumPy puede activar 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, verificará los requisitos de alineación y activará una copia si es necesario. Al pasar un búfer de CPU de matriz ND a NumPy, generalmente el búfer satisfará los requisitos de alineación y NumPy no necesitará crear una copia.

Las matrices ND pueden referirse a búferes colocados en dispositivos distintos de la memoria de la CPU local. En tales casos, la invocación de una función NumPy activará copias en la red o el dispositivo según sea necesario.

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

Precedencia del operador

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

x = tnp.ones([2]) + np.ones([2])
print("x = %s\nclass = %s" % (x, x.__class__))
x = ndarray<tf.Tensor([2. 2.], shape=(2,), dtype=float64)>
class = <class 'tensorflow.python.ops.numpy_ops.np_arrays.ndarray'>

TF NumPy y TensorFlow

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

tf.Tensor y matriz ND

La matriz ND es una envoltura delgada en tf.Tensor . Estos tipos se pueden convertir de forma económica entre sí sin generar copias de datos reales.

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

# Convert `tf.Tensor` to `ndarray`.
tnp_x = tnp.asarray(x)
print(tnp_x)

# Convert `ndarray` to `tf.Tensor` can be done in following ways.
print(tnp_x.data)
print(tf.convert_to_tensor(tnp_x))

# Note that tf.Tensor.numpy() will continue to return `np.ndarray`.
print(x.numpy(), x.numpy().__class__)
ndarray<tf.Tensor([1 2], shape=(2,), dtype=int32)>
tf.Tensor([1 2], shape=(2,), dtype=int32)
tf.Tensor([1 2], shape=(2,), dtype=int32)
[1 2] <class 'numpy.ndarray'>

Interoperabilidad de TensorFlow

Una matriz ND se puede pasar a las API de TensorFlow. Estas llamadas convierten internamente las entradas de la matriz ND en tf.Tensor . Como se mencionó anteriormente, dicha conversión en realidad no realiza copias de datos, incluso para los datos colocados en aceleradores o dispositivos remotos.

Por el contrario, los objetos tf.Tensor se pueden pasar a tf.experimental.numpy API tf.experimental.numpy . Estas entradas se convertirán internamente en una matriz ND sin realizar copias de datos.

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

# `tf.Tensor` passed into TensorFlow NumPy function.
# This returns an ND array.
tnp_sum = tnp.sum(tf.ones([2, 3]))
print("Output = %s" % tnp_sum)
Output = tf.Tensor(6.0, shape=(), dtype=float32)
Output = ndarray<tf.Tensor(6.0, shape=(), dtype=float32)>

Precedencia del operador

Cuando la matriz ND y los objetos tf.Tensor se combinan usando operadores, se usa una regla de precedencia para determinar qué objeto ejecuta el operador. Esto está controlado por el valor __array_priority__ definido por estas clases.

tf.Tensor define un __array_priority__ mayor que el de la matriz ND. Esto significa que la entrada de la matriz ND se convertirá a tf.Tensor y se tf.Tensor versión tf.Tensor del operador.

El siguiente código demuestra cómo afecta eso al tipo de salida.

x = tnp.ones([2, 2]) + tf.ones([2, 1])
print("x = %s\nClass = %s" % (x, x.__class__))
x = tf.Tensor(
[[2. 2.]
 [2. 2.]], shape=(2, 2), dtype=float32)
Class = <class 'tensorflow.python.framework.ops.EagerTensor'>

Gradientes y jacobianos: tf.GradientTape

GradientTape de TensorFlow se puede usar para la propagación hacia atrás a través de TensorFlow y el código TensorFlow NumPy. Las API de GradientTape también pueden devolver salidas de matriz ND.

Utilice 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)
Parameter shapes: [(32, 64), (64,), (64, 2)]
Gradient shapes: [(32, 64), (64,), (64, 2)]

# 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)
Output shape: (16, 2), input shape: (16, 32)
Batch jacobian shape: (16, 2, 32)

Compilación de seguimiento: tf.function

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

tf.function se puede utilizar para optimizar el código 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, y las entradas y salidas son matrices ND.

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")
Eager performance
1.7211115999998583 ms

tf.function compiled performance
0.8105368999849816 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 el elemento de lote de entrada correspondiente. Dicho cálculo se puede realizar de manera eficiente usando 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))
Weight shape: (32, 64), batch size: 128, per example gradient shape: (128, 32, 64) 
Weight shape: (64,), batch size: 128, per example gradient shape: (128, 64) 
Weight shape: (64, 2), batch size: 128, per example gradient shape: (128, 64, 2) 

# 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
    output = compute_gradients(model,
                               tnp.expand_dims(inp, 0),
                               tnp.expand_dims(label, 0))
    return output

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

print("Running vectorized computaton")
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=5) * 200, "ms")
Running vectorized computaton
0.8167428999968251 ms

Running unvectorized computation
10.86823519999598 ms

Colocación del dispositivo

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

TensorFlow también tiene API para replicar el cálculo en todos los dispositivos y realizar reducciones colectivas que no se tratarán 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"
d885545 cama

Colocación de operaciones: tf.device

Las operaciones se pueden colocar en un dispositivo llamándolo 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.data.device)
Using device: /device:CPU:0
prediction is placed on /job:localhost/replica:0/task:0/device:CPU:0

Copia de matrices ND entre dispositivos: tnp.copy

Una llamada a tnp.copy , colocada en un determinado alcance de dispositivo, copiará los datos en 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.data.device)
print(prediction_cpu.data.device)
/job:localhost/replica:0/task:0/device:CPU:0
/job:localhost/replica:0/task:0/device:CPU:0

Comparaciones de desempeño

TensorFlow NumPy usa kernels 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 traduce 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 gastos generales más altos para el envío de operaciones 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 punto de referencia a continuación 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)

png

Otras lecturas