Trang này được dịch bởi Cloud Translation API.
Switch to English

API NumPy trên TensorFlow

Xem trên TensorFlow.org Chạy trong Google Colab Xem nguồn trên GitHub Tải xuống sổ tay

Tổng quat

TensorFlow triển khai một tập hợp con của API NumPy , có sẵn dưới dạng tf.experimental.numpy . Điều này cho phép chạy mã NumPy, được tăng tốc bởi TensorFlow, đồng thời cho phép truy cập vào tất cả các API của TensorFlow.

Thiết lập

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

Mảng TensorFlow NumPy ND

Một ví dụ của tf.experimental.numpy.ndarray , được gọi là ND Array , đại diện cho một mảng dày đặc đa chiều của một loại dtype nhất định được đặt trên một thiết bị nhất định. Mỗi một trong những đối tượng này bao bọc bên trong một tf.Tensor . Kiểm tra lớp mảng ND để biết các phương thức hữu ích như ndarray.T , ndarray.reshape , ndarray.ravel và các phương thức khác.

Đầu tiên hãy tạo một đối tượng mảng ND, sau đó gọi các phương thức khác nhau.

# 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

Loại khuyến mãi

Các API TensorFlow NumPy có ngữ nghĩa được xác định rõ ràng để chuyển đổi các từ sang mảng ND, cũng như để thực hiện quảng bá kiểu trên đầu vào mảng ND. Vui lòng xem np.result_type để biết thêm chi tiết. . Khi chuyển đổi các chữ sang mảng ND, NumPy thích các loại rộng như tnp.int64tnp.float64 .

Ngược lại, tf.convert_to_tensor thích tf.int32tf.float32 để chuyển đổi hằng số thành tf.Tensor . Các API TensorFlow tf.Tensor nguyên đầu vào tf.Tensor và không thực hiện quảng bá kiểu trên chúng.

Trong ví dụ tiếp theo, bạn sẽ thực hiện quảng cáo loại. Đầu tiên, chạy bổ sung trên các đầu vào mảng ND của các loại khác nhau và lưu ý các loại đầu ra. Không có quảng cáo loại nào trong số này được phép trên các đối tượng tf.Tensor thẳng. Cuối cùng, chuyển đổi các chữ sang mảng ND bằng ndarray.asarrray và lưu ý kiểu kết quả.

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


Phát thanh truyền hình

Tương tự như TensorFlow, NumPy xác định ngữ nghĩa phong phú cho các giá trị "phát sóng". Bạn có thể xem hướng dẫn phát sóng NumPy để biết thêm thông tin và so sánh điều này với ngữ nghĩa phát sóng 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)

Lập chỉ mục

NumPy xác định các quy tắc lập chỉ mục rất phức tạp. Xem hướng dẫn Lập chỉ mục NumPy . Lưu ý việc sử dụng mảng ND như các chỉ số bên dưới.

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.

Mô hình mẫu

Tiếp theo, bạn có thể xem cách tạo một mô hình và chạy suy luận trên đó. Mô hình đơn giản này áp dụng một lớp relu theo sau là một phép chiếu tuyến tính. Các phần sau sẽ trình bày cách tính toán độ dốc cho mô hình này bằng cách sử dụng GradientTape của 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 và NumPy

TensorFlow NumPy triển khai một tập hợp con của thông số NumPy đầy đủ. Mặc dù nhiều biểu tượng sẽ được thêm vào theo thời gian, nhưng có những tính năng có hệ thống sẽ không được hỗ trợ trong tương lai gần. Chúng bao gồm hỗ trợ NumPy C API, tích hợp Swig, thứ tự lưu trữ Fortran, chế độ xem và stride_tricks , và một số loại dtype (như np.recarraynp.object ). Để biết thêm chi tiết, vui lòng xem Tài liệu API TensorFlow NumPy .

Khả năng tương tác NumPy

Mảng TensorFlow ND có thể tương tác với các hàm NumPy. Các đối tượng này triển khai giao diện __array__ . NumPy sử dụng giao diện này để chuyển đổi các đối số của hàm thành các giá trị np.ndarray trước khi xử lý chúng.

Tương tự, các hàm TensorFlow NumPy có thể chấp nhận các đầu vào thuộc nhiều loại khác nhau bao gồm tf.Tensornp.ndarray . Các đầu vào này được chuyển đổi thành mảng ND bằng cách gọi ndarray.asarray trên chúng.

Việc chuyển đổi mảng ND sang và từ np.ndarray có thể kích hoạt các bản sao dữ liệu thực tế. Vui lòng xem phần về bản sao đệm để biết thêm chi tiết.

# 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

Bản sao đệm

Việc trộn TensorFlow NumPy với mã NumPy có thể kích hoạt các bản sao dữ liệu. Điều này là do TensorFlow NumPy có các yêu cầu nghiêm ngặt hơn về căn chỉnh bộ nhớ so với NumPy.

Khi np.ndarray được chuyển tới TensorFlow Numpy, nó sẽ kiểm tra các yêu cầu về căn chỉnh và kích hoạt một bản sao nếu cần. Khi chuyển bộ đệm CPU mảng ND cho NumPy, nói chung bộ đệm sẽ đáp ứng các yêu cầu căn chỉnh và NumPy sẽ không cần tạo bản sao.

Mảng ND có thể tham chiếu đến bộ đệm được đặt trên các thiết bị khác với bộ nhớ CPU cục bộ. Trong những trường hợp như vậy, việc gọi hàm NumPy sẽ kích hoạt các bản sao trên mạng hoặc thiết bị khi cần thiết.

Do đó, việc trộn lẫn các lệnh gọi NumPy API thường phải được thực hiện một cách thận trọng và người dùng nên đề phòng chi phí sao chép dữ liệu. Việc xen kẽ các cuộc gọi TensorFlow NumPy với các cuộc gọi TensorFlow nói chung là an toàn và tránh sao chép dữ liệu. Xem phần về khả năng tương tác tensorflow để biết thêm chi tiết.

Ưu tiên điều hành

TensorFlow NumPy xác định __array_priority__ cao hơn NumPy's. Điều này có nghĩa là đối với các toán tử liên quan đến cả mảng ND và np.ndarray , đầu vào trước sẽ được ưu tiên, tức là đầu vào np.ndarray sẽ được chuyển đổi thành một mảng ND và việc triển khai TensorFlow NumPy của toán tử sẽ được gọi.

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 và TensorFlow

TensorFlow NumPy được xây dựng dựa trên TensorFlow và do đó tương tác liền mạch với TensorFlow.

tf.Tensor và mảng ND

Mảng ND là một trình bao bọc mỏng trên tf.Tensor . Các loại này có thể được chuyển đổi rẻ sang một loại khác mà không gây ra các bản sao dữ liệu thực tế.

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

Khả năng tương tác TensorFlow

Một mảng ND có thể được chuyển tới các API TensorFlow. Các lệnh gọi này chuyển đổi nội bộ đầu vào mảng ND thành tf.Tensor . Như đã đề cập trước đó, việc chuyển đổi như vậy không thực sự thực hiện sao chép dữ liệu, ngay cả đối với dữ liệu được đặt trên máy gia tốc hoặc thiết bị từ xa.

Ngược lại, các đối tượng tf.Tensor có thể được chuyển tới các API tf.experimental.numpy . Các đầu vào này sẽ được chuyển đổi nội bộ thành mảng ND mà không cần thực hiện sao chép dữ liệu.

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

Ưu tiên điều hành

Khi mảng ND và đối tượng tf.Tensor được kết hợp bằng cách sử dụng toán tử, quy tắc ưu tiên được sử dụng để xác định đối tượng nào thực thi toán tử. Điều này được kiểm soát bởi giá trị __array_priority__ được xác định bởi các lớp này.

tf.Tensor xác định __array_priority__ cao hơn so với mảng ND. Điều này có nghĩa là đầu vào mảng ND sẽ được chuyển đổi thành tf.Tensor và phiên bản tf.Tensor của toán tử sẽ được gọi.

Đoạn mã dưới đây trình bày cách điều đó ảnh hưởng đến loại đầu ra.

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

Gradients và Jacobians: tf.GradientTape

GradientTape của TensorFlow có thể được sử dụng để nhân giống ngược thông qua mã TensorFlow và TensorFlow NumPy. Các API GradientTape cũng có thể trả về kết quả đầu ra mảng ND.

Sử dụng mô hình được tạo trong phần Mô hình ví dụ và tính toán độ dốc và jacobians.

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)

Biên dịch theo dõi: tf. Chức năng

Chức năng tf. của tf.function hoạt động bằng cách "biên dịch theo dõi" mã và sau đó tối ưu hóa các dấu vết này để có hiệu suất nhanh hơn nhiều. Xem phần Giới thiệu về Đồ thị và Hàm .

tf.function cũng có thể được sử dụng để tối ưu hóa mã TensorFlow NumPy. Đây là một ví dụ đơn giản để chứng minh tốc độ tăng tốc. Lưu ý rằng nội dung của mã tf.function bao gồm các lệnh gọi đến các API TensorFlow NumPy và các đầu vào và đầu ra là các mảng 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

Vectorization: tf.vectorized_map

TensorFlow có hỗ trợ sẵn có để vectơ hóa các vòng lặp song song, cho phép tăng tốc độ từ một đến hai bậc. Các tốc độ tăng tốc này có thể truy cập được qua API tf.vectorized_map và áp dụng cho cả mã TensorFlow NumPy.

Đôi khi hữu ích khi tính toán gradient của mỗi đầu ra trong một loạt wrt phần tử lô đầu vào tương ứng. Việc tính toán như vậy có thể được thực hiện hiệu quả bằng cách sử dụng tf.vectorized_map như hình dưới đây.

@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

Vị trí thiết bị

TensorFlow NumPy có thể đặt các hoạt động trên CPU, GPU, TPU và các thiết bị từ xa. Nó sử dụng cơ chế TensorFlow tiêu chuẩn để đặt thiết bị. Dưới đây là một ví dụ đơn giản cho thấy cách liệt kê tất cả các thiết bị và sau đó đặt một số tính toán trên một thiết bị cụ thể.

TensorFlow cũng có các API để sao chép tính toán trên các thiết bị và thực hiện giảm thiểu tập thể sẽ không được đề cập ở đây.

Liệt kê các thiết bị

tf.config.list_logical_devicestf.config.list_physical_devices có thể được sử dụng để tìm thiết bị sẽ sử dụng.

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

Đặt hoạt động: tf.device

Các hoạt động có thể được thực hiện trên một thiết bị bằng cách gọi nó trong phạm vi 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

Sao chép mảng ND trên các thiết bị: tnp.copy

Một lệnh gọi tới tnp.copy , được đặt trong một phạm vi thiết bị nhất định, sẽ sao chép dữ liệu vào thiết bị đó, trừ khi dữ liệu đã có trên thiết bị đó.

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

So sánh hiệu suất

TensorFlow NumPy sử dụng hạt nhân TensorFlow được tối ưu hóa cao có thể được gửi trên CPU, GPU và TPU. TensorFlow cũng thực hiện nhiều tối ưu hóa trình biên dịch, như kết hợp hoạt động, chuyển sang cải thiện hiệu suất và bộ nhớ. Xem Tối ưu hóa đồ thị TensorFlow với Grappler để tìm hiểu thêm.

Tuy nhiên TensorFlow có chi phí cao hơn cho các hoạt động điều phối so với NumPy. Đối với khối lượng công việc bao gồm các hoạt động nhỏ (dưới khoảng 10 micro giây), các chi phí này có thể chiếm ưu thế trong thời gian chạy và NumPy có thể cung cấp hiệu suất tốt hơn. Đối với các trường hợp khác, TensorFlow nói chung sẽ cung cấp hiệu suất tốt hơn.

Chạy điểm chuẩn bên dưới để so sánh hiệu suất NumPy và TensorFlow Numpy cho các kích thước đầu vào khác nhau.

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

đọc thêm