|View source on GitHub|
Parallel map on the list of tensors unpacked from
elems on dimension 0.
Compat aliases for migration
See Migration guide for more details.
tf.vectorized_map( fn, elems, fallback_to_while_loop=True )
This method works similar to tf.map_fn but is optimized to run much faster, possibly with a much larger memory footprint. The speedups are obtained by vectorization (see https://arxiv.org/pdf/1903.04243.pdf). The idea behind vectorization is to semantically launch all the invocations of
fn in parallel and fuse corresponding operations across all these invocations. This fusion is done statically at graph generation time and the generated code is often similar in performance to a manually fused version.
tf.vectorized_map fully parallelizes the batch, this method will generally be significantly faster than using
tf.map_fn, especially in eager mode. However this is an experimental feature and currently has a lot of limitations:
fn, i.e. it should be safe to map the elements of the inputs in any order.
fnhas limited support for control flow operations.
fnshould return nested structure of Tensors or Operations. However if an Operation is returned, it should have zero outputs.
fnshould not depend on the input to
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.shape == (batch_size, num_features, 1) assert per_example_gradients.shape == (batch_size, 1)
| || The callable to be performed. It accepts one argument, which will have the same (possibly nested) structure as |
| || A tensor or (possibly nested) sequence of tensors, each of which will be unpacked along their first dimension. The nested sequence of the resulting slices will be mapped over by |
| || If true, on failing to vectorize an operation, the unsupported op is wrapped in a tf.while_loop to execute the map iterations. Note that this fallback only happens for unsupported ops and other parts of |
|A tensor or (possibly nested) sequence of tensors. Each tensor packs the results of applying fn to tensors unpacked from elems along the first dimension, from first to last.|
| ||If vectorization fails and fallback_to_while_loop is False.|
© 2020 The TensorFlow Authors. All rights reserved.
Licensed under the Creative Commons Attribution License 3.0.
Code samples licensed under the Apache 2.0 License.