Cette page a été traduite par l'API Cloud Translation.
Switch to English

API NumPy sur TensorFlow

Voir sur TensorFlow.org Exécuter dans Google Colab Afficher la source sur GitHub Télécharger le carnet

Aperçu

TensorFlow implémente un sous-ensemble de l' API NumPy , disponible en tant que tf.experimental.numpy . Cela permet d'exécuter du code NumPy, accéléré par TensorFlow, tout en permettant également l'accès à toutes les API de TensorFlow.

Installer

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

Baie TensorFlow NumPy ND

Une instance de tf.experimental.numpy.ndarray , appelée ND Array , représente un tableau dense multidimensionnel d'un dtype donné placé sur un certain périphérique. Chacun de ces objets enveloppe en interne un tf.Tensor . Consultez la classe de tableau ND pour des méthodes utiles telles que ndarray.T , ndarray.reshape , ndarray.ravel et autres.

Créez d'abord un objet de tableau ND, puis appelez différentes méthodes.

# 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

Promotion de type

Les API TensorFlow NumPy ont une sémantique bien définie pour la conversion de littéraux en tableau ND, ainsi que pour effectuer la promotion de type sur les entrées de tableau ND. Veuillez consulter np.result_type pour plus de détails. . Lors de la conversion de littéraux en tableau ND, NumPy préfère les types larges comme tnp.int64 et tnp.float64 .

En revanche, tf.convert_to_tensor préfère les types tf.int32 et tf.float32 pour convertir les constantes en tf.Tensor . Les API TensorFlow laissent les entrées tf.Tensor inchangées et n'effectuent pas de promotion de type sur celles-ci.

Dans l'exemple suivant, vous effectuerez une promotion de type. Tout d'abord, exécutez l'addition sur les entrées du tableau ND de types différents et notez les types de sortie. Aucune de ces promotions de type ne serait autorisée sur les objets tf.Tensor . Enfin, convertissez les littéraux en tableau ND à l'aide de ndarray.asarrray et notez le type résultant.

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


Diffusion

Semblable à TensorFlow, NumPy définit une sémantique riche pour les valeurs de «diffusion». Vous pouvez consulter le guide de diffusion NumPy pour plus d'informations et comparer cela avec la sémantique de diffusion 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)

Indexage

NumPy définit des règles d'indexation très sophistiquées. Consultez le guide d'indexation NumPy . Notez l'utilisation de tableaux ND comme indices ci-dessous.

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.

Exemple de modèle

Ensuite, vous pouvez voir comment créer un modèle et y exécuter des inférences. Ce modèle simple applique une couche relu suivie d'une projection linéaire. Les sections suivantes montreront comment calculer des dégradés pour ce modèle à l'aide du 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 et NumPy

TensorFlow NumPy implémente un sous-ensemble de la spécification NumPy complète. Alors que d'autres symboles seront ajoutés au fil du temps, il existe des fonctionnalités systématiques qui ne seront pas prises en charge dans un proche avenir. Il s'agit notamment de la prise en charge de l'API NumPy C, de l'intégration Swig, de l'ordre de stockage Fortran, des vues et des stride_tricks , et de certains dtype (comme np.recarray et np.object ). Pour plus de détails, consultez la documentation de l'API NumPy TensorFlow .

Interopérabilité NumPy

Les baies TensorFlow ND peuvent interagir avec les fonctions NumPy. Ces objets implémentent l'interface __array__ . NumPy utilise cette interface pour convertir les arguments de np.ndarray valeurs np.ndarray avant de les traiter.

De même, les fonctions TensorFlow NumPy peuvent accepter des entrées de différents types, notamment tf.Tensor et np.ndarray . Ces entrées sont converties en un tableau ND en appelant ndarray.asarray sur elles.

La conversion du tableau ND vers et depuis np.ndarray peut déclencher des copies de données réelles. Veuillez consulter la section sur les copies tampons pour plus de détails.

# 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

Copies tampons

Mélanger TensorFlow NumPy avec le code NumPy peut déclencher des copies de données. En effet, TensorFlow NumPy a des exigences plus strictes en matière d'alignement de la mémoire que celles de NumPy.

Lorsqu'un np.ndarray est passé à TensorFlow Numpy, il vérifie les exigences d'alignement et déclenche une copie si nécessaire. Lors du passage d'un tampon de processeur ND array à NumPy, généralement le tampon satisfera aux exigences d'alignement et NumPy n'aura pas besoin de créer une copie.

Les tableaux ND peuvent faire référence à des tampons placés sur des périphériques autres que la mémoire CPU locale. Dans de tels cas, l'appel d'une fonction NumPy déclenchera des copies sur le réseau ou le périphérique selon les besoins.

Compte tenu de cela, le mélange avec les appels d'API NumPy doit généralement être effectué avec prudence et l'utilisateur doit faire attention aux frais généraux liés à la copie de données. L'entrelacement des appels TensorFlow NumPy avec les appels TensorFlow est généralement sûr et évite de copier des données. Voir la section sur l' interopérabilité du tensorflow pour plus de détails.

Priorité des opérateurs

TensorFlow NumPy définit un __array_priority__ supérieur à NumPy's. Cela signifie que pour les opérateurs impliquant à la fois le tableau ND et np.ndarray , le premier sera prioritaire, c'est-à-dire np.ndarray entrée np.ndarray sera convertie en un tableau ND et que l'implémentation TensorFlow NumPy de l'opérateur sera appelée.

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 et TensorFlow

TensorFlow NumPy est construit sur TensorFlow et interagit donc de manière transparente avec TensorFlow.

tf.Tensor et ND

Le tableau ND est un wrapper léger sur tf.Tensor . Ces types peuvent être convertis à moindre coût les uns aux autres sans déclencher de copies de données réelles.

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'>

Interopérabilité TensorFlow

Un tableau ND peut être transmis aux API TensorFlow. Ces appels convertissent en interne les entrées du tableau ND en tf.Tensor . Comme mentionné précédemment, une telle conversion ne fait pas réellement de copies de données, même pour les données placées sur des accélérateurs ou des périphériques distants.

Inversement, les objets tf.Tensor peuvent être passés aux API tf.experimental.numpy . Ces entrées seront converties en interne en un tableau ND sans effectuer de copies de données.

# 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)>

Priorité des opérateurs

Lorsque le tableau ND et les objets tf.Tensor sont combinés à l'aide d'opérateurs, une règle de priorité est utilisée pour déterminer quel objet exécute l'opérateur. Ceci est contrôlé par la valeur __array_priority__ définie par ces classes.

tf.Tensor définit un __array_priority__ supérieur à celui du tableau ND. Cela signifie que l'entrée du tableau ND sera convertie en tf.Tensor et la version tf.Tensor de l'opérateur sera appelée.

Le code ci-dessous montre comment cela affecte le type de sortie.

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'>

Dégradés et jacobiens: tf.GradientTape

GradientTape de TensorFlow peut être utilisé pour la rétropropagation via le code NumPy TensorFlow et TensorFlow. Les API GradientTape peuvent également renvoyer des sorties de tableau ND.

Utilisez le modèle créé dans la section Exemple de modèle et calculez les dégradés et les jacobiens.

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)

Compilation de traces: tf.function

La fonction tf.function de tf.function fonctionne en «compilant» le code puis en optimisant ces traces pour des performances beaucoup plus rapides. Voir l' introduction aux graphiques et aux fonctions .

tf.function peut également être utilisé pour optimiser le code TensorFlow NumPy. Voici un exemple simple pour démontrer les accélérations. Notez que le corps du code tf.function inclut des appels aux API TensorFlow NumPy et que les entrées et la sortie sont des tableaux 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

Vectorisation: tf.vectorized_map

TensorFlow a un support intégré pour la vectorisation de boucles parallèles, ce qui permet des accélérations d'un à deux ordres de grandeur. Ces accélérations sont accessibles via l'API tf.vectorized_map et s'appliquent également au code TensorFlow NumPy.

Il est parfois utile de calculer le gradient de chaque sortie dans un lot par rapport à l'élément de lot d'entrée correspondant. Un tel calcul peut être effectué efficacement en utilisant tf.vectorized_map comme indiqué ci-dessous.

@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

Placement de l'appareil

TensorFlow NumPy peut effectuer des opérations sur des processeurs, des GPU, des TPU et des périphériques distants. Il utilise des mécanismes TensorFlow standard pour le placement des appareils. Ci-dessous, un exemple simple montre comment répertorier tous les périphériques, puis placer des calculs sur un périphérique particulier.

TensorFlow dispose également d'API pour répliquer le calcul sur plusieurs appareils et effectuer des réductions collectives qui ne seront pas couvertes ici.

Lister les appareils

tf.config.list_logical_devices et tf.config.list_physical_devices peuvent être utilisés pour trouver les périphériques à utiliser.

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"
All logical devices: [LogicalDevice(name='/device:CPU:0', device_type='CPU')]
All physical devices: [PhysicalDevice(name='/physical_device:CPU:0', device_type='CPU')]

Opérations de placement: tf.device

Les opérations peuvent être placées sur un périphérique en l'appelant dans une portée 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

Copie de baies ND sur plusieurs périphériques: tnp.copy

Un appel à tnp.copy , placé dans une certaine étendue d'appareil, copiera les données sur cet appareil, à moins que les données ne soient déjà sur cet appareil.

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

Comparaisons de performances

TensorFlow NumPy utilise des noyaux TensorFlow hautement optimisés qui peuvent être distribués sur des processeurs, des GPU et des TPU. TensorFlow effectue également de nombreuses optimisations du compilateur, comme la fusion d'opérations, qui se traduisent par des améliorations de performances et de mémoire. Voir Optimisation des graphes TensorFlow avec Grappler pour en savoir plus.

Cependant, TensorFlow a des frais généraux plus élevés pour les opérations de répartition par rapport à NumPy. Pour les charges de travail composées de petites opérations (moins d'environ 10 microsecondes), ces frais généraux peuvent dominer l'exécution et NumPy pourrait fournir de meilleures performances. Dans les autres cas, TensorFlow devrait généralement fournir de meilleures performances.

Exécutez l'analyse comparative ci-dessous pour comparer les performances de NumPy et TensorFlow Numpy pour différentes tailles d'entrée.

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

Lectures complémentaires