Question

首先，我知道我应该使用top_k，但是使k-max池难以实现（在TF中实现）的原因是必须保留顺序。

我到目前为止有什么：

import tensorflow as tf
from tensorflow.contrib.framework import sort
sess = tf.Session()
a = tf.convert_to_tensor([[[5, 1, 10, 2], [3, 11, 2, 6]]])
b = sort(tf.nn.top_k(a, k=2)[1])
print(tf.gather(a, b, axis=-1).eval(session=sess))

已经很近了，但是还没到

我得到的是：[[[[[[5，5]，[1，2]]]，[[[3，2]，[11，6]]]]]

我想要什么：[[[5，10]，[11，6]]

我几乎百分百地确定是否需要collect_nd，但我不知道这是什么，而且我是pytorch用户，在那里真的很容易

import torch
a = torch.LongTensor([[[5, 1, 10, 2], [3, 11, 2, 6]]])
b = a.topk(2, dim = -1)[1].sort(dim = -1)[0]
print(a.gather(-1, b))

哦，还有我发现的每个代码都不是一个保存命令的命令（从语义上来说是错误的）

Answer 1

很奇怪...这应该超级简单，但是我们找不到合适的解决方案...

尝试一下：

sess = tf.Session()
k = 2
a = tf.convert_to_tensor([[[5, 1, 10, 2], [3, 11, 2, 6]]])
b = tf.nn.top_k(a, k=k, sorted=True)[1]
b = sort(b)

flatA = tf.reshape(a,(-1,))
shapeA = tf.shape(a)
lenA = tf.shape(flatA)[0]
kShape = tf.concat([shapeA[:-1],tf.constant([k])], axis=-1)

indices = tf.range(lenA)
indices = tf.reshape(indices,shapeA)

toSum = tf.expand_dims(tf.gather(indices,0,axis=-1), axis=-1)
b += toSum
b = tf.reshape(b,(-1,))

gat = tf.gather(flatA, b)
gat = tf.reshape(gat, kShape)

print(gat.eval(session=sess))

Answer 2

尝试使用这个Custom层，将上面的代码修改为一个自定义的tensorflow keras层。

class KMaxPooling(layers.Layer):
"""
K-max pooling layer that extracts the k-highest activations from a sequence (2nd dimension).
TensorFlow backend.
"""
def __init__(self, k=1, axis=1, **kwargs):
    super(KMaxPooling, self).__init__(**kwargs)
    self.input_spec = layers.InputSpec(ndim=3)
    self.k = k

    assert axis in [1,2],  'expected dimensions (samples, filters, convolved_values),\
               cannot fold along samples dimension or axis not in list [1,2]'
    self.axis = axis

    # need to switch the axis with the last elemnet
    # to perform transpose for tok k elements since top_k works in last axis
    self.transpose_perm = [0,1,2] #default
    self.transpose_perm[self.axis] = 2
    self.transpose_perm[2] = self.axis

def compute_output_shape(self, input_shape):
    input_shape_list = list(input_shape)
    input_shape_list[self.axis] = self.k
    return tuple(input_shape_list)

def call(self, x):
    # swap sequence dimension to get top k elements along axis=1
    transposed_for_topk = tf.transpose(x, perm=self.transpose_perm)

    # extract top_k, returns two tensors [values, indices]
    top_k_vals, top_k_indices = tf.math.top_k(transposed_for_topk,
                                              k=self.k, sorted=True,
                                              name=None)
    # maintain the order of values as in the paper
    # sort indices
    sorted_top_k_ind = tf.sort(top_k_indices)
    flatten_seq = tf.reshape(transposed_for_topk, (-1,))
    shape_seq = tf.shape(transposed_for_topk)
    len_seq = tf.shape(flatten_seq)[0]
    indices_seq = tf.range(len_seq)
    indices_seq = tf.reshape(indices_seq, shape_seq)
    indices_gather = tf.gather(indices_seq, 0, axis=-1)
    indices_sum = tf.expand_dims(indices_gather, axis=-1)
    sorted_top_k_ind += indices_sum
    k_max_out = tf.gather(flatten_seq, sorted_top_k_ind)
    # return back to normal dimension but now sequence dimension has only k elements
    # performing another transpose will get the tensor back to its original shape
    # but will have k as its axis_1 size
    transposed_back = tf.transpose(k_max_out, perm=self.transpose_perm)

    return transposed_back

如何在Tensorflow或Keras中实现K-Max池化？

2 个答案: