使用 tf.gather 或 tf.gather_nd答案

【问题标题】：Using tf.gather or tf.gather_nd使用 tf.gather 或 tf.gather_nd
【发布时间】：2021-03-03 12:17:39
【问题描述】：

我有一个维度为(BATCH_SIZE*A*B*FEATURE_LENGTH) 的输入。现在我想从每个输入样本的每个 A 块中选择 k(out of B) 行。每个 A 块的 k 值是不同的。例如。

inp = ([[[[ 5, 38, 40, 13, 28],
         [12,  6, 36, 20, 23],
         [44, 35, 23, 46,  3]],

        [[22, 32, 36, 20, 42],
         [ 0, 19, 41, 36, 17],
         [ 9, 35, 44,  7, 19]],

        [[27, 10, 22, 10, 48],
         [16, 42, 27,  7, 38],
         [35, 32, 15, 39, 28]]]])
#size (1,3,3,5) = (1,A,B,FEATURE_LENGTH)

现在说 k=2，即我想从 3 个块中的每个块中提取 2 行。我想要

row 0 and 1 from 1st block
row 1 and 2 from 2nd block
row 0 and 2 from 3rd block

这意味着我希望我的输出看起来像

([[[[ 5, 38, 40, 13, 28],
    [12,  6, 36, 20, 23]],
   
    [[ 0, 19, 41, 36, 17],
     [ 9, 35, 44,  7, 19]],

    [[27, 10, 22, 10, 48],
     [35, 32, 15, 39, 28]]]])
#op shape = (1,3,2,5)

我发现如果我们提供索引作为，使用tf.gather_nd 这是可能的

ind = array([[[[0, 0, 0], [0, 0, 1]], [[0, 1, 1], [0, 1, 2]], [[0, 2, 0], [0, 2, 2]]]])

但如果我输入大小为(1,16,16,128) 和k=4，创建这个长索引序列将变得乏味。在 Tensorflow-2 中有更简单的方法吗？谢谢！

【问题讨论】：

标签： python tensorflow deep-learning tensorflow2.0

【解决方案1】：

使用tf.gather() 和batch_dims 参数：

inds = tf.constant([[[0, 1], [1, 2], [0, 2]]])
output = tf.gather(inp, inds, batch_dims=2)

【讨论】：

谢谢！那行得通..我在问这里之前尝试使用 tf.gather 但是把 axis=2 而不是 batch_dims=2 这导致了问题！