如何在不重复元素的情况下从张量中选择top-n个元素？

潜行：

鉴于拾取的元素都是唯一的，我想选择3维张量的top-n个元素。所有元素都按第二列排序，在下面的示例中，我选择了top-2，但我不想在其中重复。

条件：否for loops或tf.map_fn()
这是我想要的输入和期望的输出：

input_tensor = tf.constant([
  [[2.0, 1.0],
   [2.0, 1.0],
   [3.0, 0.4],
   [1.0, 0.1]],
  [[44.0, 0.8],
   [22.0, 0.7],
   [11.0, 0.5],
   [11.0, 0.5]],
  [[5555.0, 0.8],
   [3333.0, 0.7],
   [4444.0, 0.4],
   [1111.0, 0.1]],
  [[444.0, 0.8],
   [333.0, 1.1],
   [333.0, 1.1],
   [111.0, 0.1]]
])

这就是我现在要得到的。我不要！

>> TOPK = 2
>> topk_resutls = tf.gather(
    input_tensor,
    tf.math.top_k(input_tensor[:, :, 1], k=TOPK, sorted=True).indices,
    batch_dims=1
)
>> topk_resutls.numpy().tolist()
[[[2.0, 1.0], [2.0, 1.0]],
 [[44.0, 0.8], [22.0, 0.7]],
 [[5555.0, 0.8], [3333.0, 0.7]],
 [[333.0, 1.1], [333.0, 1.1]]]

这是我真正想要的

[[[2.0, 1.0], [3.0, 0.4]],       # [3.0, 0.4] is the 2nd highest element based on 2nd column
 [[44.0, 0.8], [22.0, 0.7]],   
 [[5555.0, 0.8], [3333.0, 0.7]],
 [[333.0, 1.1], [444.0, 0.8]]]   # [444.0, 0.8] is the 2nd highest element based on 2nd column

jdehesa：

这是一种可行的方法，尽管由于它首先对数组进行排序，因此需要做更多的工作。

import tensorflow as tf
import numpy as np

# Input data
k = 2
input_tensor = tf.constant([
  [[2.0, 1.0],
   [2.0, 1.0],
   [3.0, 0.4],
   [1.0, 0.1]],
  [[44.0, 0.8],
   [22.0, 0.7],
   [11.0, 0.5],
   [11.0, 0.5]],
  [[5555.0, 0.8],
   [3333.0, 0.7],
   [4444.0, 0.4],
   [1111.0, 0.1]],
  [[444.0, 0.8],
   [333.0, 1.1],
   [333.0, 1.1],
   [111.0, 0.1]]
])
# Sort by first column
idx = tf.argsort(input_tensor[..., 0], axis=-1)
s = tf.gather_nd(input_tensor, tf.expand_dims(idx, axis=-1), batch_dims=1)
# Find repeated elements
col1 = s[..., 0]
col1_ext = tf.concat([col1[..., :1] - 1, col1], axis=-1)
mask = tf.math.not_equal(col1_ext[..., 1:], col1_ext[..., :-1])
# Replace value for repeated elements with "minus infinity"
col2 = s[..., 1]
col2_masked = tf.where(mask, col2, col2.dtype.min)
# Get top-k results
topk_idx = tf.math.top_k(col2_masked, k=k, sorted=True).indices
topk_results = tf.gather(s, topk_idx, batch_dims=1)
# Print
with np.printoptions(suppress=True):
    print(topk_results.numpy())
# [[[   2.     1. ]
#   [   3.     0.4]]
# 
#  [[  44.     0.8]
#   [  22.     0.7]]
# 
#  [[5555.     0.8]
#   [3333.     0.7]]
# 
#  [[ 333.     1.1]
#   [ 444.     0.8]]]

请注意，存在一种极端情况，即k组中没有不同的元素。在这种情况下，此解决方案会将重复的元素放在末尾，这会破坏得分顺序。

本文收集自互联网，转载请注明来源。

如有侵权，请联系 [email protected] 删除。