tf.vectorized_map

Параллельное применение функции к списку тензоров, распакованных из elems по размерности 0.

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

Используется в блокнотах

Этот метод работает аналогично методу tf.map_fn, но оптимизирован для значительно большей скорости, возможно, с гораздо большим объёмом занимаемой памяти. Ускорение достигается за счёт векторизации (см. Автоматическая векторизация графиков TensorFlow: Якобианы, автоматическое пакетное выполнение и не только). Идея векторизации заключается в семантическом параллельном запуске всех вызовов fn и слиянии соответствующих операций по всем этим вызовам. Это слияние выполняется статически во время генерации графика, и полученный код часто демонстрирует производительность, сопоставимую с вручную объединённой версией.

Поскольку tf.vectorized_map полностью паралелизует пакет, этот метод, как правило, будет значительно быстрее, чем использование tf.map_fn, особенно в режиме eager. Однако это экспериментальная функция, и в настоящее время она имеет много ограничений:

• Между различными семантическими вызовами fn не должно быть зависимости данных, то есть должно быть безопасно сопоставлять элементы входных данных в любом порядке.
• Состоятельные ядра могут в основном не поддерживаться, поскольку они часто подразумевают зависимость данных. Тем не менее, мы поддерживаем ограниченный набор таких состоятельных ядер (например, RandomFoo, операции с переменными, такие как чтение и т. д.).
• fn имеет ограниченную поддержку операций управления потоком.
• fn должен возвращать вложенную структуру тензоров или операций. Однако если возвращается операция, у неё должно быть нулевое количество выходов.
• Форма и тип данных любых промежуточных или выходных тензоров в вычислении fn не должны зависеть от входных данных для 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)