TensorFlow：tf.map_fn函数

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

map_fn 的最简单版本反复地将可调用的 fn 应用于从第一个到最后一个的元素序列.这些元素由 elems 解压缩的张量构成.dtype 是 fn 的返回值的数据类型.如果与elems 的数据类型不同,用户必须提供 dtype.

这种方法也允许 fn 的多元 elems 和输出

,如果 elems 是：(t1, [t2, t3, [t4, t5]]),则 fn 的适当签名为：fn = lambda (t1, [t2, t3, [t4, t5]]):.

例子

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]