Zapisz datę! Google I / O powraca w dniach 18-20 maja Zarejestruj się teraz
Ta strona została przetłumaczona przez Cloud Translation API.
Switch to English

NumPy API w TensorFlow

Zobacz na TensorFlow.org Uruchom w Google Colab Wyświetl źródło w serwisie GitHub Pobierz notatnik

Przegląd

TensorFlow implementuje podzbiór NumPy API , dostępny jako tf.experimental.numpy . Pozwala to na uruchamianie kodu NumPy przyspieszonego przez TensorFlow, a jednocześnie umożliwia dostęp do wszystkich interfejsów API TensorFlow.

Ustawiać

pip install -q tf-nightly
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.5.0-dev20210225

Włączam zachowanie NumPy

Aby użyć tnp jako NumPy, włącz zachowanie NumPy dla TensorFlow:

tnp.experimental_enable_numpy_behavior()

To wywołanie umożliwia promocję typu w TensorFlow, a także zmienia wnioskowanie o typie podczas konwertowania literałów na tensory, aby bardziej ściśle przestrzegać standardu NumPy.

Tablica TensorFlow NumPy ND

Instancja tf.experimental.numpy.ndarray , zwana ND Array , reprezentuje wielowymiarową gęstą tablicę danego typu dtype umieszczoną na określonym urządzeniu. Jest to alias do tf.Tensor . Sprawdź klasę tablicową ND, aby uzyskać przydatne metody, takie jak ndarray.T , ndarray.reshape , ndarray.ravel i inne.

Najpierw utwórz obiekt tablicy ND, a następnie wywołaj różne metody.

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

Is `ones` an instance of tf.Tensor: True

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

Wpisz promocję

Interfejsy API TensorFlow NumPy mają dobrze zdefiniowaną semantykę do konwersji literałów na tablicę ND, a także do wykonywania promocji typów na wejściach tablic ND. Aby uzyskać więcej informacji, zobacz np.result_type .

Interfejsy API TensorFlow pozostawiają tf.Tensor wejściowe tf.Tensor niezmienione i nie wykonują na nich promocji typów, podczas gdy interfejsy API TensorFlow NumPy promują wszystkie dane wejściowe zgodnie z regułami promocji typu NumPy. W następnym przykładzie przeprowadzisz promocję typu. Najpierw uruchom dodatek na wejściach tablicy ND różnych typów i zanotuj typy danych wyjściowych. Żadna z promocji tego typu nie byłaby dozwolona przez interfejsy API TensorFlow.

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))
Type promotion for operations
int32 + int64 => int64
int32 + float32 => float64
int32 + float64 => float64
int64 + float32 => float64
int64 + float64 => float64
float32 + float64 => float64

Na koniec przekonwertuj literały na tablicę ND za pomocą ndarray.asarray i zanotuj wynikowy typ.

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)
Type inference during array creation
tnp.asarray(1).dtype == tnp.int64
tnp.asarray(1.).dtype == tnp.float64

Konwertując literały na tablicę ND, NumPy preferuje szerokie typy, takie jak tnp.int64 i tnp.float64 . Natomiast tf.convert_to_tensor preferuje tf.int32 i tf.float32 do konwersji stałych do tf.Tensor . Interfejsy API TensorFlow NumPy są zgodne z zachowaniem NumPy dla liczb całkowitych. Jeśli chodzi o elementy prefer_float32 argument prefer_float32 parametru experimental_enable_numpy_behavior pozwala kontrolować, czy preferować tf.float32 tf.float64 (domyślnie False ). Na przykład:

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)
When prefer_float32 is True:
tnp.asarray(1.).dtype == tnp.float32
tnp.add(1., 2.).dtype == tnp.float32
When prefer_float32 is False:
tnp.asarray(1.).dtype == tnp.float64
tnp.add(1., 2.).dtype == tnp.float64

Nadawanie

Podobnie jak TensorFlow, NumPy definiuje bogatą semantykę dla wartości „rozgłaszania”. Aby uzyskać więcej informacji, zapoznaj się z przewodnikiem dotyczącym transmisji NumPy i porównaj je z semantyką nadawania 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)

Indeksowanie

NumPy definiuje bardzo wyrafinowane reguły indeksowania. Zobacz przewodnik po indeksowaniu NumPy . Zwróć uwagę na użycie tablic ND jako indeksów poniżej.

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
tf.Tensor(
[[[16 17 18 19]
  [20 21 22 23]]], shape=(1, 2, 4), dtype=int64) 

Boolean indexing
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
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.

Przykładowy model

Następnie możesz zobaczyć, jak utworzyć model i przeprowadzić na nim wnioskowanie. Ten prosty model nakłada warstwę relu, po której następuje rzutowanie liniowe. Późniejsze sekcje pokażą, jak obliczyć gradienty dla tego modelu przy użyciu GradientTape 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)))
tf.Tensor(
[[ 0.6435826  -0.14300033]
 [ 0.6435826  -0.14300033]], shape=(2, 2), dtype=float32)

TensorFlow NumPy i NumPy

TensorFlow NumPy implementuje podzbiór pełnej specyfikacji NumPy. Chociaż z czasem będzie dodawanych więcej symboli, istnieją systematyczne funkcje, które nie będą obsługiwane w najbliższej przyszłości. Należą NumPy C wsparcie API, integracja Swig, Fortran porządek składowania, widoki i stride_tricks , a niektóre dtype s (jak np.recarray i np.object ). Więcej informacji można znaleźć w dokumentacji TensorFlow NumPy API .

Współdziałanie NumPy

Macierze TensorFlow ND mogą współpracować z funkcjami NumPy. Te obiekty implementują interfejs __array__ . NumPy używa tego interfejsu do konwersji argumentów funkcji na wartości np.ndarray przed ich przetworzeniem.

Podobnie funkcje TensorFlow NumPy mogą akceptować dane wejściowe różnych typów, w tym np.ndarray . Te dane wejściowe są konwertowane na tablicę ND przez wywołanie na nich ndarray.asarray .

Konwersja tablicy ND do iz np.ndarray może wyzwolić rzeczywiste kopie danych. Więcej informacji można znaleźć w sekcji dotyczącej kopii w buforze .

# 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.framework.ops.EagerTensor'>
# It is easy to plot ND arrays, given the __array__ interface.
labels = 15 + 2 * tnp.random.randn(1, 1000)
_ = plt.hist(labels)

png

Kopie buforowe

Mieszanie TensorFlow NumPy z kodem NumPy może powodować kopiowanie danych. Dzieje się tak, ponieważ TensorFlow NumPy ma bardziej rygorystyczne wymagania dotyczące wyrównania pamięci niż NumPy.

Gdy np.ndarray zostanie przekazany do TensorFlow NumPy, sprawdzi wymagania dotyczące wyrównania i w razie potrzeby uruchomi kopię. Podczas przekazywania buforu procesora tablicy ND do NumPy, zwykle bufor spełnia wymagania dotyczące wyrównania, a NumPy nie będzie musiał tworzyć kopii.

Macierze ND mogą odnosić się do buforów umieszczonych na urządzeniach innych niż lokalna pamięć procesora. W takich przypadkach wywołanie funkcji NumPy wyzwoli kopie w sieci lub urządzeniu, jeśli zajdzie taka potrzeba.

Biorąc to pod uwagę, mieszanie z wywołaniami NumPy API powinno być generalnie wykonywane ostrożnie, a użytkownik powinien uważać na narzuty związane z kopiowaniem danych. Przeplatanie wywołań TensorFlow NumPy z wywołaniami TensorFlow jest ogólnie bezpieczne i pozwala uniknąć kopiowania danych. Więcej informacji można znaleźć w sekcji poświęconej współdziałaniu TensorFlow .

Pierwszeństwo operatorów

TensorFlow NumPy definiuje __array_priority__ wyższy niż NumPy. Oznacza to, że dla operatorów obejmujących zarówno tablicę ND, jak i np.ndarray , pierwszeństwo będzie miała ta pierwsza, tj. Wejście np.ndarray zostanie przekonwertowane na tablicę ND i wywołana zostanie implementacja operatora TensorFlow NumPy.

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

TF NumPy i TensorFlow

TensorFlow NumPy jest oparty na TensorFlow i dzięki temu bezproblemowo współpracuje z TensorFlow.

tf.Tensor i tablica ND

Tablica ND jest aliasem do tf.Tensor , więc oczywiście można je mieszać bez wyzwalania rzeczywistych kopii danych.

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

Współdziałanie TensorFlow

Tablicę ND można przekazać do interfejsów API TensorFlow, ponieważ tablica ND jest tylko aliasem do tf.Tensor . Jak wspomniano wcześniej, taka współpraca nie powoduje kopiowania danych, nawet w przypadku danych umieszczonych na akceleratorach lub zdalnych urządzeniach.

I odwrotnie, obiekty tf.Tensor mogą być przekazywane do API tf.experimental.numpy bez wykonywania kopii danych.

# 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)
Output = tf.Tensor(6.0, shape=(), dtype=float32)
Output = tf.Tensor(6.0, shape=(), dtype=float32)

Gradienty i jakobiany: tf.GradientTape

GradientTape TensorFlow może być używany do wstecznej propagacji poprzez kod TensorFlow i TensorFlow NumPy.

Użyj modelu utworzonego w sekcji Przykładowy model i oblicz gradienty oraz jakobiany.

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

Kompilacja śledzenia: tf.function

Funkcja tf.function tf.function działa poprzez „śledzenie kompilacji” kodu, a następnie optymalizację tych śladów w celu uzyskania znacznie większej wydajności. Zobacz wprowadzenie do wykresów i funkcji .

tf.function może również służyć do optymalizacji kodu TensorFlow NumPy. Oto prosty przykład pokazujący przyspieszenia. Zwróć uwagę, że tf.function kodu tf.function zawiera wywołania funkcji API 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")
Eager performance
1.6580017000023872 ms

tf.function compiled performance
0.41147700000010445 ms

Wektoryzacja: tf.vectorized_map

TensorFlow ma wbudowaną obsługę wektoryzacji równoległych pętli, co umożliwia przyspieszenie od jednego do dwóch rzędów wielkości. Te przyspieszenia są dostępne za pośrednictwem interfejsu API tf.vectorized_map i mają również zastosowanie do kodu TensorFlow NumPy.

Czasami przydatne jest obliczenie gradientu każdego wyjścia w partii z odpowiednim wejściowym elementem wsadowym. Takie obliczenia można przeprowadzić wydajnie za pomocą tf.vectorized_map jak pokazano poniżej.

@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
    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")
Running vectorized computation
0.38309449999474054 ms

Running unvectorized computation
40.52958750000357 ms

Umiejscowienie urządzenia

TensorFlow NumPy może umieszczać operacje na procesorach, układach GPU, TPU i urządzeniach zdalnych. Wykorzystuje standardowe mechanizmy TensorFlow do umieszczania urządzeń. Poniżej prosty przykład pokazuje, jak wyświetlić listę wszystkich urządzeń, a następnie umieścić obliczenia na określonym urządzeniu.

TensorFlow ma również interfejsy API do replikowania obliczeń na różnych urządzeniach i wykonywania zbiorczych redukcji, które nie zostaną tutaj omówione.

Lista urządzeń

tf.config.list_logical_devices i tf.config.list_physical_devices mogą służyć do znalezienia używanych urządzeń.

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

Operacje umieszczania: tf.device

Operacje można umieszczać na urządzeniu, wywołując je w zakresie 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)
Using device: LogicalDevice(name='/device:GPU:0', device_type='GPU')
prediction is placed on /job:localhost/replica:0/task:0/device:GPU:0

Kopiowanie tablic ND między urządzeniami: tnp.copy

Wywołanie tnp.copy , umieszczone w określonym zakresie urządzenia, spowoduje skopiowanie danych do tego urządzenia, chyba że dane już znajdują się na tym urządzeniu.

with tf.device("/device:CPU:0"):
  prediction_cpu = tnp.copy(prediction)
print(prediction.device)
print(prediction_cpu.device)
/job:localhost/replica:0/task:0/device:GPU:0
/job:localhost/replica:0/task:0/device:CPU:0

Porównania wydajności

TensorFlow NumPy wykorzystuje wysoce zoptymalizowane jądra TensorFlow, które mogą być wysyłane na CPU, GPU i TPU. TensorFlow wykonuje również wiele optymalizacji kompilatorów, takich jak fuzja operacji, co przekłada się na poprawę wydajności i pamięci. Zobacz Optymalizacja wykresów TensorFlow za pomocą Grapplera, aby dowiedzieć się więcej.

Jednak TensorFlow ma większe narzuty na operacje wysyłania w porównaniu z NumPy. W przypadku obciążeń składających się z małych operacji (mniej niż około 10 mikrosekund) te narzuty mogą dominować w czasie wykonywania, a NumPy może zapewnić lepszą wydajność. W innych przypadkach TensorFlow powinien generalnie zapewniać lepszą wydajność.

Uruchom poniższy test porównawczy, aby porównać wydajność NumPy i TensorFlow NumPy dla różnych rozmiarów wejściowych.

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

Dalsza lektura