Python 从tf.nn.max\u pool\u使用argmax tensorflow

我正在尝试使用

tf.nn.max\u pool\u的argmax结果和\u argmax（）

来索引另一个张量。为了简单起见，假设我正在尝试实现以下内容：

output, argmax = tf.nn.max_pool_with_argmax(input, ksize, strides, padding)
tf.assert_equal(input[argmax],output)

现在我的问题是如何实现必要的索引操作

input[argmax]

，以获得所需的结果？我猜这涉及到

tf.gather\n（）

和相关调用的一些用法，但我无法理解。如有必要，我们可以假设输入具有

[BatchSize，Height，Width，Channel]

维度

谢谢你的帮助

---

Mat

我找到了一个使用tf.gather\n和的解决方案，它很有效，尽管看起来不那么优雅。我使用了发布的函数

unravel\u argmax

如果有人有更优雅的解决方案，我还是很想知道

---

Mat

我是这样做的：

def max_pool(input, ksize, strides,padding):
    output, arg_max = tf.nn.max_pool_with_argmax(input=input,ksize=ksize,strides=strides,padding=padding)
    shape=tf.shape(output)
    output1=tf.reshape(tf.gather(tf.reshape(input,[-1]),arg_max),shape)

    err=tf.reduce_sum(tf.square(tf.subtract(output,output1)))
    return output1, err

---

这个小片段可以工作：

def get_results(data,other_tensor):
    pooled_data, indices = tf.nn.max_pool_with_argmax(data,ksize=[1,ksize,ksize,1],strides=[1,stride,stride,1],padding='VALID',include_batch_in_index=True)
    b,w,h,c = other_tensor.get_shape.as_list()
    other_tensor_pooled = tf.gather(tf.reshape(other_tensor,shape= [b*w*h*c,]),indices)
    return other_tensor_pooled

---

上述

索引

可用于索引张量。此函数实际上返回展平索引，要将其用于

batch\u size>1

的任何对象，您需要将

include\u batch\u作为
True
传递到\u index

，以获得正确的结果。我在这里假设

othertensor

您的批处理大小与

数据相同。

您是否检查了输入[argmax]本身不起作用？Tensorflow支持相对高级的跨步切片（如numpy），因此这可能会起作用（尽管您可能需要将其应用于填充输入）。这不起作用…必须在_index=True中包含参数

include_batch_

，才能在多个图像上正常工作

def get_results(data,other_tensor):
    pooled_data, indices = tf.nn.max_pool_with_argmax(data,ksize=[1,ksize,ksize,1],strides=[1,stride,stride,1],padding='VALID',include_batch_in_index=True)
    b,w,h,c = other_tensor.get_shape.as_list()
    other_tensor_pooled = tf.gather(tf.reshape(other_tensor,shape= [b*w*h*c,]),indices)
    return other_tensor_pooled