TensorFlow中的高阶函数:tf.map_fn()
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]
以免误操作。