Python tf.vectorized_map用法及代码示例

从维度 0 的 elems 解压缩的张量列表上的并行映射。

用法

tf.vectorized_map(
    fn, elems, fallback_to_while_loop=True
)

参数

fn 要执行的可调用对象。它接受一个参数，该参数将具有与 elems 相同的(可能是嵌套的)结构，并返回一个可能嵌套的张量和操作结构，该结构可能与 elems 的结构不同。
elems 张量或(可能是嵌套的)张量序列，每个张量都将沿其第一维展开。结果切片的嵌套序列将由 fn 映射。所有元素的第一个维度必须广播到一致的值；等效地，对于某些常见的批量大小 B >= 1 ，每个元素张量必须具有 B 或 1 的第一维。
fallback_to_while_loop 如果为 true，则在未能矢量化操作时，将不支持的操作包装在 tf.while_loop 中以执行映射迭代。请注意，这种回退只发生在不受支持的操作上，fn 的其他部分仍然是矢量化的。如果为 false，则在遇到不受支持的操作时，会抛出 ValueError。请注意，回退可能会导致减速，因为矢量化通常会产生一到两个数量级的加速。

张量或(可能是嵌套的)张量序列。每个张量从第一维到最后一个维度，将 fn 应用于从 elems 解包的张量的结果打包。
尽管它们作为 user-visible 输入和输出不太常见，但请注意，表示张量列表的类型为 tf.variant 的张量(例如来自 tf.raw_ops.TensorListFromTensor )是通过堆叠列表内容而不是变量本身来矢量化的，因此容器张量在返回时将具有标量形状，而不是通常的堆叠形状。这提高了控制流梯度矢量化的性能。

抛出

ValueError 如果矢量化失败并且 fallback_to_while_loop 为 False。

此方法的用法方式类似于tf.map_fn，但经过优化以运行得更快，可能具有更大的内存占用。加速是通过矢量化获得的(参见Auto-Vectorizing TensorFlow Graphs:Jacobians, Auto-Batching and Beyond)。矢量化背后的想法是在语义上并行启动 fn 的所有调用，并在所有这些调用中融合相应的操作。这种融合是在图形生成时静态完成的，生成的代码在性能上通常与手动融合的版本相似。

因为 tf.vectorized_map 完全并行化批处理，所以此方法通常比使用 tf.map_fn 快得多，尤其是在即刻模式下。然而，这是一个实验性函数，目前有很多限制：

fn 的不同语义调用之间不应存在数据依赖性，即以任何顺序映射输入的元素应该是安全的。
有状态内核可能大多不受支持，因为它们通常意味着数据依赖。不过，我们确实支持一组有限的此类有状态内核(如 RandomFoo、读取等变量操作等)。
fn 对控制流操作的支持有限。
fn 应该返回张量或操作的嵌套结构。但是，如果返回一个操作，它应该有零输出。
fn 计算中任何中间或输出张量的形状和 dtype 不应依赖于 fn 的输入。

例子：

def outer_product(a):
  return tf.tensordot(a, a, 0)

batch_size = 100
a = tf.ones((batch_size, 32, 32))
c = tf.vectorized_map(outer_product, a)
assert c.shape == (batch_size, 32, 32, 32, 32)

# Computing per-example gradients

batch_size = 10
num_features = 32
layer = tf.keras.layers.Dense(1)

def model_fn(arg):
  with tf.GradientTape() as g:
    inp, label = arg
    inp = tf.expand_dims(inp, 0)
    label = tf.expand_dims(label, 0)
    prediction = layer(inp)
    loss = tf.nn.l2_loss(label - prediction)
  return g.gradient(loss, (layer.kernel, layer.bias))

inputs = tf.random.uniform([batch_size, num_features])
labels = tf.random.uniform([batch_size, 1])
per_example_gradients = tf.vectorized_map(model_fn, (inputs, labels))
assert per_example_gradients[0].shape == (batch_size, num_features, 1)
assert per_example_gradients[1].shape == (batch_size, 1)

相关用法

注：本文由纯净天空筛选整理自tensorflow.org大神的英文原创作品 tf.vectorized_map。非经特殊声明，原始代码版权归原作者所有，本译文未经允许或授权，请勿转载或复制。

用法

参数

返回

抛出

例子：