このページは Cloud Translation API によって翻訳されました。
Switch to English

TensorFlowのNumPyAPI

TensorFlow.orgで表示 GoogleColabで実行 GitHubでソースを表示ノートブックをダウンロード

概要概要

TensorFlow用具のサブセットnumpyののAPIとして利用でき、 tf.experimental.numpy 。これにより、TensorFlowによって高速化されたNumPyコードを実行できると同時に、TensorFlowのすべてのAPIにアクセスできます。

セットアップ

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

TensorFlow NumPyNDアレイ

ND Arrayと呼ばれるtf.experimental.numpy.ndarrayのインスタンスは、特定のデバイスに配置された特定のdtype多次元の密な配列を表します。これらのオブジェクトはそれぞれ、内部でtf.Tensorラップしtf.Tensorndarray.Tndarray.reshapendarray.ravelなどの便利なメソッドについては、ND配列クラスを確認してください。

最初にND配列オブジェクトを作成してから、さまざまなメソッドを呼び出します。

# 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

タイププロモーション

TensorFlow NumPy APIには、リテラルをND配列に変換するため、およびND配列入力で型プロモーションを実行するための明確に定義されたセマンティクスがあります。詳細については、 np.result_typeを参照してください。 。リテラルをND配列に変換する場合、NumPyはtnp.int64tnp.float64などのワイドタイプをtnp.int64tnp.float64

対照的に、 tf.convert_to_tensorは、定数をtf.Tensor変換するためにtf.int32およびtf.float32タイプをtf.Tensorます。 TensorFlow APIは、 tf.Tensor入力を変更せず、タイププロモーションを実行しません。

次の例では、タイププロモーションを実行します。まず、さまざまなタイプのND配列入力で加算を実行し、出力タイプをメモします。これらのタイプのプロモーションはいずれも、ストレートtf.Tensorオブジェクトでは許可されません。最後に、 ndarray.asarrrayを使用してリテラルをND配列に変換し、結果の型をメモします。

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


放送

TensorFlowと同様に、NumPyは「ブロードキャスト」値の豊富なセマンティクスを定義します。詳細については、 NumPyブロードキャストガイドを確認し、これを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)

インデックス作成

NumPyは、非常に洗練されたインデックス作成ルールを定義しています。 NumPyインデックスガイドを参照してください。以下のインデックスとしてND配列を使用していることに注意してください。

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.

モデルの例

次に、モデルを作成して推論を実行する方法を確認できます。この単純なモデルは、reluレイヤーとそれに続く線形射影を適用します。後のセクションでは、TensorFlowのGradientTapeを使用してこのモデルの勾配を計算する方法を示します。

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

TensorFlowNumPyとNumPy

TensorFlow NumPyは、完全なNumPy仕様のサブセットを実装します。時間の経過とともにシンボルが追加される予定ですが、近い将来サポートされなくなる体系的な機能があります。これらには、NumPy C APIサポート、 stride_tricks統合、Fortranストレージ順序、ビューとstride_tricks 、およびいくつかのdtypenp.recarraynp.object )が含まれます。詳細については、 TensorFlow NumPyAPIドキュメントをご覧ください。

NumPyの相互運用性

TensorFlow ND配列は、NumPy関数と相互運用できます。これらのオブジェクトは、 __array__インターフェイスを実装します。 NumPyはこのインターフェースを使用して、関数の引数を処理する前にnp.ndarray値に変換します。

同様に、TensorFlow NumPy関数は、 tf.Tensornp.ndarrayなどのさまざまなタイプの入力を受け入れることができます。これらの入力は、それらに対してndarray.asarrayを呼び出すことにより、ND配列に変換されます。

ND配列をnp.ndarrayとの間で変換すると、実際のデータコピーがトリガーされる場合があります。詳細については、バッファコピーのセクションを参照してください。

# 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

バッファコピー

TensorFlow NumPyをNumPyコードと混在させると、データコピーがトリガーされる場合があります。これは、TensorFlowNumPyのメモリアライメントに関する要件がNumPyの要件よりも厳しいためです。

np.ndarraynp.ndarrayに渡されると、アライメント要件を確認し、必要に応じてコピーをトリガーします。 NDアレイCPUバッファーをNumPyに渡す場合、通常、バッファーはアライメント要件を満たし、NumPyはコピーを作成する必要はありません。

ND配列は、ローカルCPUメモリ以外のデバイスに配置されたバッファを参照できます。このような場合、NumPy関数を呼び出すと、必要に応じてネットワークまたはデバイス全体でコピーがトリガーされます。

このため、NumPy API呼び出しとの混合は通常、注意して行う必要があり、ユーザーはデータのコピーのオーバーヘッドに注意する必要があります。 TensorFlow NumPy呼び出しとTensorFlow呼び出しのインターリーブは一般的に安全であり、データのコピーを回避します。詳細については、テンソルフローの相互運用性に関するセクションを参照してください。

演算子の優先順位

TensorFlow NumPyは、NumPyよりも高い__array_priority__定義します。これは、ND配列とnp.ndarray両方を含む演算子の場合、前者が優先されることをnp.ndarrayます。つまり、 np.ndarray入力がND配列に変換され、演算子のTensorFlowNumPy実装が呼び出されます。

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

TFNumPyとTensorFlow

TensorFlow NumPyはTensorFlowの上に構築されているため、TensorFlowとシームレスに相互運用できます。

tf.TensorおよびND配列

ND配列は、 tf.Tensor薄いラッパーtf.Tensor 。これらのタイプは、実際のデータコピーをトリガーすることなく、安価に相互に変換できます。

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

TensorFlowの相互運用性

ND配列をTensorFlowAPIに渡すことができます。これらの呼び出しは、ND配列入力を内部的にtf.Tensor変換しtf.Tensor 。前述のように、このような変換では、アクセラレータやリモートデバイスに配置されたデータであっても、実際にはデータのコピーは行われません。

逆に、 tf.Tensorオブジェクトがに渡すことができtf.experimental.numpyのAPI。これらの入力は、データコピーを実行せずに、内部でND配列に変換されます。

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

演算子の優先順位

ND配列とtf.Tensorオブジェクトが演算子を使用して結合される場合、優先ルールを使用して、どのオブジェクトが演算子を実行するかを決定します。これは、これらのクラスによって定義された__array_priority__値によって制御されます。

tf.Tensor定義__array_priority__ NDアレイよりも高いです。これは、ND配列入力がtf.Tensorに変換され、演算子のtf.Tensorバージョンが呼び出されることを意味します。

以下のコードは、それが出力タイプにどのように影響するかを示しています。

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

グラデーションとヤコビアン:tf.GradientTape

TensorFlowのGradientTapeは、TensorFlowおよびTensorFlowNumPyコードを介したバックプロパゲーションに使用できます。 GradientTape APIは、ND配列出力を返すこともできます。

「モデルの例」セクションで作成したモデルを使用して、勾配とヤコビアンを計算します。

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)

トレースコンパイル:tf.function

tf.functiontf.functionは、コードを「トレースコンパイル」し、これらのトレースを最適化してパフォーマンスを大幅に向上させることで機能します。グラフと関数の概要を参照してください。

tf.functionを使用して、TensorFlowNumPyコードを最適化することもできます。これは、スピードアップを示す簡単な例です。 tf.functionコードの本体にはtf.function APIの呼び出しが含まれており、入力と出力は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

ベクトル化:tf.vectorized_map

TensorFlowには、並列ループのベクトル化のサポートが組み込まれているため、1〜2桁の高速化が可能です。これらのtf.vectorized_mapは、 tf.vectorized_map APIを介してアクセスでき、 tf.vectorized_mapコードにも適用されます。

対応する入力バッチ要素を使用して、バッチ内の各出力の勾配を計算すると便利な場合があります。このような計算は、以下に示すようにtf.vectorized_mapを使用して効率的に実行できます。

@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

デバイスの配置

TensorFlow NumPyは、CPU、GPU、TPU、およびリモートデバイスに操作を配置できます。デバイスの配置に標準のTensorFlowメカニズムを使用します。以下の簡単な例は、すべてのデバイスを一覧表示してから、特定のデバイスに計算を配置する方法を示しています。

TensorFlowには、デバイス間で計算を複製し、ここでは取り上げない集合的な削減を実行するためのAPIもあります。

デバイスの一覧表示

tf.config.list_logical_devicesおよびtf.config.list_physical_devicesを使用して、使用するデバイスを見つけることができます。

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

配置操作: tf.device

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

デバイス間でのNDアレイのコピー: tnp.copy

特定のデバイススコープに配置されたtnp.copy呼び出すと、データがそのデバイスに既に存在しない限り、そのデバイスにデータがコピーされます。

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

パフォーマンスの比較

TensorFlow NumPyは、CPU、GPU、TPUにディスパッチできる高度に最適化されたTensorFlowカーネルを使用します。 TensorFlowは、操作の融合など、パフォーマンスとメモリの改善につながる多くのコンパイラ最適化も実行します。詳細については、Grapplerを使用したTensorFlowグラフの最適化をご覧ください。

ただし、TensorFlowは、NumPyと比較してディスパッチ操作のオーバーヘッドが高くなります。小さな操作(約10マイクロ秒未満)で構成されるワークロードの場合、これらのオーバーヘッドがランタイムを支配する可能性があり、NumPyはより優れたパフォーマンスを提供する可能性があります。その他の場合、TensorFlowは一般的にパフォーマンスが向上するはずです。

以下のベンチマークを実行して、さまざまな入力サイズでのNumPyとTensorFlowNumpyのパフォーマンスを比較します。

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

参考文献