TensorFlow中的高阶函数：tf.map_fn()

$\nabla$ 联系方式：

e-mail: FesianXu@gmail.com

QQ: 973926198

github: https://github.com/FesianXu

知乎专栏: 计算机视觉/计算机图形理论与应用

微信公众号：

---

TensorFlow中的高阶函数：tf.map_fn()

在TensorFlow中，有一些函数被称为高阶函数（high-level function），和在python中的高阶函数意义相似，其也是将函数当成参数传入，以实现一些有趣的，有用的操作。其中tf.map_fn()就是其中一个。我们这里介绍一下这个函数。

首先引入一个TF在应用上的问题：一般我们处理图片的时候，常常用到卷积，也就是tf.nn.conv2d()，但是这个函数的输入格式如：(batch_size, image_height, image_width, image_channel)，其中batch_size为一个批次的大小，我们注意到，如果按照这种输入的话，我们只能对一张图片进行卷积操作。在需要对视频进行卷积操作的时候，因为视频的输入格式一般如：(batch_size, video_size, frame_height, frame_width, frame_channel)，其中多出了一个视频的长度video_size，这样我们就不能简单对视频进行卷积了。我们用一张图描述我们面临的问题和map_fn是如何帮我们解决的：

其中，我们的数据就是从clip #1到clip #n，我们需要对每一个clip应用相同的方法function，比如是卷积函数tf.nn.conv2d()，从而得到batch_size个结果result #1到result #n。这个就是map_fn()的效果，而且，因为各个batch之间没有关联，所以可以并行快速高效地处理。我们再看看这个函数的用法，先贴出其API手册：

tf.map_fn

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

其中的fn是一个可调用的(callable)函数，就是我们图中的function，一般会使用lambda表达式表示。elems是需要做处理的Tensors，TF将会将elems从第一维展开，进行map处理。主要就是那么两个，其中dtype为可选项，但是比较重要，他表示的是fn函数的输出类型，如果fn返回的类型和elems中的不同，那么就必须显式指定为和fn返回类型相同的类型。下面给出一个视频卷积的例子：

batch = data.get_batch()  # batch with the shape of (batch_size, video_size, frame_height, frame_width, frame_channel)
cnn = tf.map_fn(fn=lambda inp: tf.nn.conv2d(inp, kernel, stride, padding='SAME'),
                elems=batch,
                dtype=tf.float32)

这样我们就对每一个batch的（batch-wise）进行了卷积处理了，大家不妨试试吧！

Update 2020/3/18:
评论区有朋友指出：

这里是不是应该将输入的shape 转变成 – > [video_size,batch_size,w,h,c]

我想在这里进行讨论下，tf.map_fn()是对输入的第一维进行展开的，因此如果你输入的张量的shape是[video_size, batch_size, width, height, channel]， 那么map_fn()会对每一个样本[batch_size, width, height, channel]进行处理，比如卷积等等，那么输出也应该是[batch_size, width_transformed, height_transformed, channel_transformed]， 这样当然是可以的，但是如果你输入的shape是[batch_size, video_size, width, height, channel]，那么每次处理的就是[video_size, w,h,c]了对吧，但是因为我们一般的处理，比如图片的二维卷积，通常都是只考虑单帧的处理，也即是一个视频中的视频帧与视频帧之间是没有关系的，那么可以设想，在这种假设下，其实你的第一维和第二维的顺序是batch_size或者video_size其实是没有关系的，可以互换，因为不管怎么样，map_fn()都只是对单帧进行处理。

但是需要注意的是，有时候我们这个假设不成立，我们需要对整个视频进行比如conv3D之类的处理，这个时候我们就必须采用[batch_size, video_size, w,h,c]了，因为如果你反过来了，就会变成在batch的方向进行3D卷积了，这不是很荒谬吗，哈哈哈，应该是在单个视频的时序方向进行3D卷积才对吧。

综上所述，我还是建议输入的shape保持[batch_size, video_size, w,h,c]以免误操作。