TensorFlow2.4 tensorflow.experim
TensorFlow 实现了 NumPy API 的子集，可作为 tf.experimental.numpy 获得。这允许运行由 TensorFlow 加速的 NumPy 代码，同时还允许访问所有 TensorFlow 的API。
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__)
TensorFlow NumPy ND array
tf.experimental.numpy.ndarray 的实例称为 ND Array，表示放置在特定设备上的给定 dtype 的多维密集数组。它是 tf.Tensor 的别名。请查看 ND数组类，以获取有用的方法，例如 ndarray.T，ndarray.reshape，ndarray.ravel等。
首先创建一个 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 " % (ones.shape, ones.ndim, ones.dtype))
# `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)
Type promotion
TensorFlow NumPy API具有明确定义的语义，用于将 字面量 转换为 ND 数组，以及对 ND 数组输入执行类型提升。有关更多详细信息，请参见np.result_type。
TensorFlow API保持 tf.Tensor 输入不变并且不对其执行类型提升，而TensorFlow NumPy API 根据 NumPy 类型提升规则来提升所有输入。在下一个示例中，您将执行类型提升。首先，对不同类型的ND阵列输入运行加法并记下输出类型。TensorFlow API 不允许这些类型的提升。
最后，使用 ndarray.asarray 将 字面量 转换为 ND数组，并记下结果类型。
Broadcasting
与 TensorFlow 类似，NumPy 为“广播”值定义了丰富的语义。您可以查看 NumPy broadcasting guide 以获取更多信息，并将其与 TensorFlow broadcasting semantics 进行比较。
Indexing
NumPy 定义了非常复杂的索引规则。请参阅 NumPy Indexing guide。请注意，下面将 ND 数组用作索引。
Example Model
接下来，您将看到如何创建模型并对其进行推理。这个简单的模型应用 relu 层，然后进行线性投影。后面的部分将展示如何使用 TensorFlow 的 GradientTape 计算该模型的梯度。
class Model:
"""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(
[[1.477266 0.53620607]
[1.477266 0.53620607]], shape=(2, 2), dtype=float32)>
TensorFlow NumPy and NumPy
TensorFlow NumPy 实现了完整的 NumPy 规范的子集。虽然随着时间的推移会添加更多符号，但是有些系统功能将在不久的将来不再支持。其中包括对 NumPy C API 的支持，Swig 集成，Fortran 存储顺序，视图和stride_tricks 以及一些 dtype（例如 np.recarray 和 np.object）。有关更多详细信息，请参阅 TensorFlow NumPy API Documentation。
NumPy 互操作性
TensorFlow ND arrays 可以与 NumPy 函数互操作。这些对象实现__array__ 接口。NumPy 使用此接口在处理函数参数之前将其转换为 np.ndarray 值。
同样，TensorFlow NumPy 函数可以接受不同类型的输入，包括 np.ndarray。通过在它们上调用 ndarray.asarray 将这些输入转换为 ND数组。
ND 数组与 np.ndarray 之间的转换可能会触发实际的数据副本。有关更多详细信息，请参见 buffer copies 部分。
Buffer copies
将 TensorFlow NumPy 与 NumPy 代码混合在一起可能会触发数据复制。这是因为 TensorFlow NumPy 对内存对齐的要求比 NumPy 严格。
当 np.ndarray 传递给 TensorFlow NumPy 时，它将检查对齐要求，并在需要时触发副本。将 ND array CPU 缓冲区传递给 NumPy 时，通常该缓冲区将满足对齐要求，并且 NumPy 无需创建副本。
ND array 可以引用放置在本地 CPU 内存以外的设备上的缓冲区。在这种情况下，调用 NumPy 函数将根据需要触发网络或设备上的副本。
鉴于此，通常应谨慎地与 NumPy API 调用混合使用，并且用户应注意复制数据的开销。将 TensorFlow NumPy 调用与 TensorFlow 调用进行交错通常是安全的，并且避免了复制数据。有关更多详细信息，请参见有关 TensorFlow interoperability 的部分。
Operator precedence
TensorFlow NumPy 定义了一个比 NumPy 高的 __array_priority__。这意味着对于同时涉及 ND 数组和 np.ndarray 的运算符，前者将优先处理，即 np.ndarray 输入将转换为 ND 数组，并且该运算符的 TensorFlow NumPy 实现将被调用。
TF NumPy and TensorFlow
TensorFlow NumPy 构建在 TensorFlow 之上，因此可与 TensorFlow 无缝互操作。
tf.Tensor 和 ND array
ND数组是 tf.Tensor 的别名，因此显然可以在不触发实际数据副本的情况下将它们混合在一起。
TensorFlow互操作性
ND 数组可以传递给 TensorFlow API，因为 ND 数组只是 tf.Tensor 的别名。如前所述，即使是放置在加速器或远程设备上的数据，这种互操作也不会进行数据复制。
相反，可以将 tf.Tensor 对象传递给 tf.experimental.numpy API，而无需执行数据复制。
Gradients and Jacobians: tf.GradientTape
TensorFlow 的 GradientTape 可通过 TensorFlow 和 TensorFlow NumPy 代码用于反向传播。
使用在 Example Model 部分中创建的模型，并计算梯度和雅各布矩阵。
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)
Trace compilation: tf.function
TensorFlow 的 tf.function 通过“跟踪编译”代码，然后优化这些跟踪以提高性能来工作。请参见 Introduction to Graphs and Functions。
tf.function 也可以用于优化 TensorFlow NumPy 代码。这是一个演示加速的简单示例。请注意，tf.function 代码的主体包括对 TensorFlow NumPy API 的调用。
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.268819999995685 ms
tf.function compiled performance
0.730070000008709 ms
Vectorization: tf.vectorized_map
TensorFlow 具有对并行循环进行矢量化的内置支持，可将速度提高一到两个数量级。这些加速可以通过 tf.vectorized_map API 访问，也适用于 TensorFlow NumPy 代码。
有时以批量计算每个输出的梯度非常有用。相应的输入批处理元素。如下所示，可以使用 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
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
1.006039999992936 ms
Running unvectorized computation
5.103150000013557 ms
Device placement
TensorFlow NumPy 可以在 CPU，GPU，TPU 和远程设备上进行操作。它使用标准的 TensorFlow 机制进行设备放置。下面的一个简单示例显示了如何列出所有设备，然后在特定设备上进行一些计算。
TensorFlow 还具有用于在设备之间复制计算并执行集体缩减的API，在此不介绍。
List devices
可以使用 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'), 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')]
Placing operations: tf.device
Operations can be placed on a device by calling it in a tf.device scope.
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
Copying ND arrays across devices: tnp.copy
A call to tnp.copy, placed in a certain device scope, will copy the data to that device, unless the data is already on that device.
with tf.device("/device:CPU:0"):
prediction_cpu = tnp.copy(prediction)
print(prediction.device)
print(prediction_cpu.device)
输出：
