tf.map_fn

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

tf.map_fn(
    fn, elems, dtype=None, parallel_iterations=None, back_prop=True,
    swap_memory=False, infer_shape=True, name=None
)

Самая простая версия map_fn последовательно применяет вызываемую функцию fn к последовательности элементов от первого до последнего. Элементы представляют собой тензоры, распакованные из elems. Тип данных возвращаемого значения fn — dtype. Пользователи должны предоставить dtype, если он отличается от типа данных elems.

Предположим, что elems распаковывается в values, список тензоров. Форма результирующего тензора — [values.shape[0]] + fn(values[0]).shape.

Этот метод также позволяет использовать многоаргументные elems и вывод fn. Если elems представляет собой (возможно, вложенный) список или кортеж тензоров, то каждый из этих тензоров должен иметь соответствующую первую (распаковываемую) размерность. Подпись fn может соответствовать структуре elems. То есть, если elems равно (t1, [t2, t3, [t4, t5]]), то соответствующая подпись для fn имеет вид: fn = lambda (t1, [t2, t3, [t4, t5]]):.

Кроме того, fn может генерировать структуру, отличную от входной. Например, fn может выглядеть так: fn = lambda t1: return (t1 + 1, t1 - 1). В этом случае параметр dtype не является необязательным: dtype должен быть типом или (возможно, вложенным) кортежем типов, соответствующим выводу fn.

Для применения функциональной операции к ненулевым элементам SparseTensor рекомендуется использовать один из следующих методов. Во-первых, если функция выражается через TensorFlow-операции, используйте

result = SparseTensor(input.indices, fn(input.values), input.dense_shape)

Если же функция не выражается через TensorFlow-операцию, используйте

result = SparseTensor(
  input.indices, map_fn(fn, input.values), input.dense_shape)

вместо этого.

При выполнении с помощью eager execution, map_fn не выполняется параллельно, даже если parallel_iterations установлено в значение > 1. Вы по-прежнему можете получить преимущества от параллельного выполнения функции, используя декоратор tf.contrib.eager.defun,

# Assume the function being used in map_fn is fn.
# To ensure map_fn calls fn in parallel, use the defun decorator.
@tf.contrib.eager.defun
def func(tensor):
  return tf.map_fn(fn, tensor)

Обратите внимание, что при использовании декоратора defun любой не относящийся к TensorFlow Python-код, написанный в вашей функции, не будет выполнен. Подробнее см. tf.contrib.eager.defun. Рекомендуется отлаживать без defun, но переключиться на defun, чтобы получить преимущества от параллельного выполнения map_fn.

Примеры:

elems = np.array([1, 2, 3, 4, 5, 6])
squares = map_fn(lambda x: x * x, elems)
# squares == [1, 4, 9, 16, 25, 36]

elems = (np.array([1, 2, 3]), np.array([-1, 1, -1]))
alternate = map_fn(lambda x: x[0] * x[1], elems, dtype=tf.int64)
# alternate == [-1, 2, -3]

elems = np.array([1, 2, 3])
alternates = map_fn(lambda x: (x, -x), elems, dtype=(tf.int64, tf.int64))
# alternates[0] == [1, 2, 3]
# alternates[1] == [-1, -2, -3]