tf.map_fn

Преобразует elems, применяя fn к каждому элементу, извлечённому по оси 0. (устаревшие аргументы)

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

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

См. также tf.scan.

map_fn извлекает elems по оси 0, получая последовательность элементов; вызывает fn для преобразования каждого элемента; затем снова объединяет преобразованные значения в тензор.

Функции отображения с одним тензором на входе и выходе

Если elems — это один тензор, а сигнатура fn имеет вид tf.Tensor->tf.Tensor, то map_fn(fn, elems) эквивалентна tf.stack([fn(elem) for elem in tf.unstack(elems)]). Например:

tf.map_fn(fn=lambda t: tf.range(t, t + 3), elems=tf.constant([3, 5, 2]))
<tf.Tensor: shape=(3, 3), dtype=int32, numpy=
  array([[3, 4, 5],
         [5, 6, 7],
         [2, 3, 4]], dtype=int32)>

map_fn(fn, elems).shape = [elems.shape[0]] + fn(elems[0]).shape.

Функции отображения с несколькими входами и выходами

map_fn также поддерживает функции с несколькими входами и выходами:

• Если elems — это кортеж (или вложенная структура) тензоров, все эти тензоры должны иметь одинаковый размер по внешнему измерению (num_elems); для преобразования каждого кортежа (или структуры) соответствующих срезов из elems используется fn. Например, если elems — это кортеж (t1, t2, t3), то fn используется для преобразования каждого кортежа срезов (t1[i], t2[i], t3[i]) (где 0 <= i < num_elems).

• Если fn возвращает кортеж (или вложенную структуру) тензоров, результат формируется путём объединения соответствующих элементов этих структур.

Указание сигнатуры выходных данных fn

Если сигнатуры входных и выходных данных fn различаются, сигнатуру выходных данных необходимо указать с помощью fn_output_signature. (Сигнатуры входных и выходных данных различаются, если не совпадают их структуры, типы данных или типы тензоров.) Например:

tf.map_fn(fn=tf.strings.length,  # input & output have different dtypes
          elems=tf.constant(["hello", "moon"]),
          fn_output_signature=tf.int32)
<tf.Tensor: shape=(2,), dtype=int32, numpy=array([5, 4], dtype=int32)>
tf.map_fn(fn=tf.strings.join,  # input & output have different structures
          elems=[tf.constant(['The', 'A']), tf.constant(['Dog', 'Cat'])],
          fn_output_signature=tf.string)
<tf.Tensor: shape=(2,), dtype=string,
 numpy=array([b'TheDog', b'ACat'], dtype=object)>

fn_output_signature можно указать одним из следующих способов:

• tf.DType или tf.TensorSpec (для описания tf.Tensor)
• tf.RaggedTensorSpec (для описания tf.RaggedTensor)
• tf.SparseTensorSpec (для описания tf.sparse.SparseTensor)
• Кортеж, список или словарь (возможно, вложенные), содержащие типы, перечисленные выше.

Нерегулярные тензоры

map_fn поддерживает входные и выходные данные типа tf.RaggedTensor. В частности:

• Если elems — это RaggedTensor, fn будет вызвана для каждой строки этого нерегулярного тензора.

  • Если elems содержит только одно нерегулярное измерение, значения, передаваемые в fn, будут объектами типа tf.Tensor.
  • Если elems содержит несколько нерегулярных измерений, значения, передаваемые в fn, будут объектами типа tf.RaggedTensor с на одно нерегулярное измерение меньше.

• Если результатом map_fn должен быть объект типа RaggedTensor, используйте tf.RaggedTensorSpec, чтобы задать fn_output_signature.

  • Если fn возвращает объекты типа tf.Tensor разных размеров, используйте tf.RaggedTensorSpec с ragged_rank=0, чтобы объединить их в один нерегулярный тензор (с ragged_rank=1).
  • Если fn возвращает объекты типа tf.RaggedTensor, используйте tf.RaggedTensorSpec с тем же ragged_rank.

# Example: RaggedTensor input
rt = tf.ragged.constant([[1, 2, 3], [], [4, 5], [6]])
tf.map_fn(tf.reduce_sum, rt, fn_output_signature=tf.int32)
<tf.Tensor: shape=(4,), dtype=int32, numpy=array([6, 0, 9, 6], dtype=int32)>

# Example: RaggedTensor output
elems = tf.constant([3, 5, 0, 2])
tf.map_fn(tf.range, elems,
          fn_output_signature=tf.RaggedTensorSpec(shape=[None],
                                                  dtype=tf.int32))
<tf.RaggedTensor [[0, 1, 2], [0, 1, 2, 3, 4], [], [0, 1]]>

Примечание: map_fn следует использовать только в том случае, если нужно применить функцию к строкам объекта типа RaggedTensor. Если нужно применить функцию к отдельным значениям, используйте:

• tf.ragged.map_flat_values(fn, rt) (если fn можно выразить с помощью операций TensorFlow)
• rt.with_flat_values(map_fn(fn, rt.flat_values)) (в противном случае)

Например:

rt = tf.ragged.constant([[1, 2, 3], [], [4, 5], [6]])
tf.ragged.map_flat_values(lambda x: x + 2, rt)
<tf.RaggedTensor [[3, 4, 5], [], [6, 7], [8]]>

Разреженные тензоры

map_fn поддерживает входные и выходные данные типа tf.sparse.SparseTensor. В частности:

• Если elems — это SparseTensor, fn будет вызвана для каждой строки этого разреженного тензора. В частности, значение, передаваемое в fn, будет объектом типа tf.sparse.SparseTensor с на одно измерение меньше, чем у elems.

• Если результатом map_fn должен быть объект типа SparseTensor, используйте tf.SparseTensorSpec, чтобы задать fn_output_signature. Отдельные объекты SparseTensor, возвращённые fn, будут объединены в один объект SparseTensor с дополнительным измерением.

# Example: SparseTensor input
st = tf.sparse.SparseTensor([[0, 0], [2, 0], [2, 1]], [2, 3, 4], [4, 4])
tf.map_fn(tf.sparse.reduce_sum, st, fn_output_signature=tf.int32)
<tf.Tensor: shape=(4,), dtype=int32, numpy=array([2, 0, 7, 0], dtype=int32)>

# Example: SparseTensor output
tf.sparse.to_dense(
    tf.map_fn(tf.sparse.eye, tf.constant([2, 3]),
              fn_output_signature=tf.SparseTensorSpec(None, tf.float32)))
<tf.Tensor: shape=(2, 3, 3), dtype=float32, numpy=
  array([[[1., 0., 0.],
          [0., 1., 0.],
          [0., 0., 0.]],
         [[1., 0., 0.],
          [0., 1., 0.],
          [0., 0., 1.]]], dtype=float32)>

Примечание: map_fn следует использовать только в том случае, если нужно применить функцию к строкам объекта типа SparseTensor. Если нужно применить функцию к ненулевым значениям, используйте:

• Если функцию можно выразить с помощью операций TensorFlow, используйте:

  tf.sparse.SparseTensor(st.indices, fn(st.values), st.dense_shape)
  

• В противном случае используйте:

  tf.sparse.SparseTensor(st.indices, tf.map_fn(fn, st.values),
                         st.dense_shape)
  

map_fn и векторизованные операции

map_fn применяет операции, используемые fn, к каждому элементу elems, выполняя в итоге O(elems.shape[0]) операций. Это отчасти компенсируется тем, что map_fn может обрабатывать элементы параллельно. Однако преобразование, выраженное с помощью map_fn, обычно всё же менее эффективно, чем эквивалентное преобразование с использованием векторизованных операций.

Обычно map_fn следует использовать только в одном из следующих случаев:

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

Например, приведённый выше пример, в котором fn=lambda t: tf.range(t, t + 3) применяется к elems, можно переписать эффективнее, используя векторизованные операции:

elems = tf.constant([3, 5, 2])
tf.range(3) + tf.expand_dims(elems, 1)
<tf.Tensor: shape=(3, 3), dtype=int32, numpy=
  array([[3, 4, 5],
         [5, 6, 7],
         [2, 3, 4]], dtype=int32)>

В некоторых случаях для автоматического преобразования функции в эквивалентную векторизованную функцию можно использовать tf.vectorized_map.

Немедленное выполнение

При немедленном выполнении map_fn не запускается параллельно, даже если для parallel_iterations задано значение > 1. Тем не менее можно повысить производительность, запустив функцию параллельно с помощью декоратора tf.function:

fn=lambda t: tf.range(t, t + 3)
@tf.function
def func(elems):
  return tf.map_fn(fn, elems, parallel_iterations=3)
func(tf.constant([3, 5, 2]))
<tf.Tensor: shape=(3, 3), dtype=int32, numpy=
  array([[3, 4, 5],
         [5, 6, 7],
         [2, 3, 4]], dtype=int32)>

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

>>> elems = np.array([1, 2, 3, 4, 5, 6])
>>> tf.map_fn(lambda x: x * x, elems)
<tf.Tensor: shape=(6,), dtype=int64, numpy=array([ 1,  4,  9, 16, 25, 36])>

elems = (np.array([1, 2, 3]), np.array([-1, 1, -1]))
tf.map_fn(lambda x: x[0] * x[1], elems, fn_output_signature=tf.int64)
<tf.Tensor: shape=(3,), dtype=int64, numpy=array([-1,  2, -3])>

elems = np.array([1, 2, 3])
tf.map_fn(lambda x: (x, -x), elems,
         fn_output_signature=(tf.int64, tf.int64))
(<tf.Tensor: shape=(3,), dtype=int64, numpy=array([1, 2, 3])>,
 <tf.Tensor: shape=(3,), dtype=int64, numpy=array([-1, -2, -3])>)