tensorflow tf.nn.max_pool_with_argmax返回最大池化对应索引值

在看Segnet的时候发现使用了带有索引的最大池化（max_pool)，在tensorflow的实现中，对应的函数为tf.nn.max_pool_with_argmax()，其返回值为所取最大值位置的索引，但采用了一种指数的计算方式进行表示

这里为官方注释

The indices in `argmax` are flattened, so that a maximum value at position
  `[b, y, x, c]` becomes flattened index
  `((b * height + y) * width + x) * channels + c`.

可以看到对索引的表示增加了一些运算操作，使得我们无法直接使用这个索引对应的位置

这里自己写了一个拆解方法的代码，记录一下（只考虑了通道影响，没有考虑batch的影响，可能后续会加以改进说明）

def unravel_corr(list,shape):
    b,h,w,c = list.get_shape().as_list()
    cc = []
    for i in range(c):
        a = tf.constant(i,dtype=tf.int64)
        a = [a]
        aa = (tf.tile(tf.tile(a,[h])[:,tf.newaxis],[1,w])[:,:,tf.newaxis])
        cc.append(aa)
    aa = tf.concat(cc,2)
    aa = tf.tile(aa[tf.newaxis,:,:,:],[b,1,1,1])
    _,height,width,_ = shape
    list_y = (list-aa)/c%width
    list_x = (list-aa)/c//width%height
    return [list_x,list_y]

整体测试代码如下

import tensorflow as tf
a = tf.constant([[5,8,2,1],[4,3,5,7],[0,7,9,1],[2,3,9,7]])
a = tf.reshape(a,[1,4,4,1])
c=tf.constant([[1,5,7,1],[4,8,5,7],[0,7,9,12],[15,3,9,7]])
c = tf.reshape(c,[1,4,4,1])
a = tf.concat([a,c],3)
b,list = tf.nn.max_pool_with_argmax(a,[1,2,2,1],[1]+[2,2]+[1],padding='VALID')
def unravel_corr(list,shape):
    b,h,w,c = list.get_shape().as_list()
    cc = []
    for i in range(c):
        a = tf.constant(i,dtype=tf.int64)
        a = [a]
        aa = (tf.tile(tf.tile(a,[h])[:,tf.newaxis],[1,w])[:,:,tf.newaxis])
        cc.append(aa)
    aa = tf.concat(cc,2)
    aa = tf.tile(aa[tf.newaxis,:,:,:],[b,1,1,1])
    _,height,width,_ = shape
    list_y = (list-aa)/c%width
    list_x = (list-aa)/c//width%height
    return [list_x,list_y]
cor = unravel_corr(list,[1,4,4,2])

with tf.Session() as sess:
    print('a:',sess.run(a))
    print('b:',sess.run(b))
    print('list:',sess.run(cor))

经计算，得到的位置索引与手动计算得到的结果相符