Python tf.map_fn用法及代码示例

通过将 fn 应用于轴 0 上未堆叠的每个元素来转换 elems。(不推荐使用的参数)

用法

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
)

警告：不推荐使用某些参数：(dtype)。它们将在未来的版本中被删除。更新说明：改用fn_output_signature

另见tf.scan。

map_fn在轴0上解栈elems以获得元素序列；调用 fn 来转换每个元素；然后将转换后的值重新堆叠在一起。

使用 single-Tensor 输入和输出映射函数

如果 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.

使用 multi-arity 输入和输出映射函数

map_fn 还支持带有 multi-arity 输入和输出的函数：

• 如果elems 是张量的元组(或嵌套结构)，那么这些张量必须都具有相同的outer-dimension 大小(num_elems)； fn 用于从 elems 转换相应切片的每个元组(或结构)。例如，如果 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 可以使用以下任何一种方式指定：

• A tf.DType 或 tf.TensorSpec(用于说明 tf.Tensor )
• A tf.RaggedTensorSpec(说明 tf.RaggedTensor )
• A tf.SparseTensorSpec(说明 tf.sparse.SparseTensor )
• 包含上述类型的(可能是嵌套的)元组、列表或字典。

不规则张量

map_fn 支持 tf.RaggedTensor 输入和输出。特别是：

• 如果 elems 是 RaggedTensor ，那么 fn 将使用该参差不齐的张量的每一行调用。

  • 如果 elems 只有一个参差不齐的维度，则传递给 fn 的值将是 tf.Tensor s。
  • 如果 elems 有多个不规则维度，则传递给 fn 的值将是少一个不规则维度的 tf.RaggedTensor 。

• 如果 map_fn 的结果应该是 RaggedTensor ，则使用 tf.RaggedTensorSpec 来指定 fn_output_signature 。

  • 如果 fn 返回具有不同大小的 tf.Tensor s，则使用 tf.RaggedTensorSpec 和 ragged_rank=0 将它们组合成单个不规则张量(将具有 ragged_rank=1)。
  • 如果 fn 返回 tf.RaggedTensor s，则使用具有相同 ragged_rank 的 tf.RaggedTensorSpec 。

# 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 ops)
• 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 的值将是一个比 elems 少一维的 tf.sparse.SparseTensor。

• 如果 map_fn 的结果应该是 SparseTensor ，则使用 tf.SparseTensorSpec 来指定 fn_output_signature 。 fn 返回的单个 SparseTensor 将堆叠成一个具有多维的单个 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 ops，请使用：

  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 可用于自动将函数转换为矢量化等效函数。

即刻执行

当即刻执行时，即使 parallel_iterations 设置为 > 1 的值，map_fn 也不会并行执行。您仍然可以获得使用 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 装饰器，您在函数中编写的任何非 TensorFlow Python 代码都不会被执行。有关详细信息，请参阅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])>)