首先,我知道我应该使用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))
哦,还有我发现的每个代码都不是一个保存命令的命令(从语义上来说是错误的)
答案 0 :(得分: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))
答案 1 :(得分:0)
尝试使用这个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