如何根據我從張量流中的另一個矩陣獲得的最大和次要值以及索引來獲取矩陣中每一行的值？

Question

如何根據我從張量流中的另一個矩陣獲得的最大和次級值索引獲取矩陣中每一行的值？ 例如我有一個矩陣張量A

    [[1,2,3],
     [6,5,4],
     [7,9,8]], 

和矩陣張量B

    [[10,11,12],
     [13,14,15],
     [16,17,18]]. 

然后我得到最大值和次要最大值索引向量形式矩陣A

    [[2,1],
     [0,2],
     [1,2]] 

通過使用tf.nn_topk。 然后我想從這些索引中獲取矩陣B的凈值，即

    [[12,11],
     [13,15],
     [17,18]]. 

我應該怎么做？ 似乎tf.gather_nd可以完成這項工作，但是我不知道如何為它提供2D索引。

Answer 1

因此，對於這種特定情況，此代碼返回值。

1. 沒有向量化
2. 它不是通用的

它只是像這樣為gather_nd創建一個模板。

 [[0 1]
  [0 2]
  [1 2]
  [1 0]
  [2 2]
  [2 1]]

其他人可能有更緊湊的想法。

import tensorflow as tf

A = tf.Variable([[10,11,12],
     [13,14,15],
     [16,17,18]], )

B = tf.Variable([[2,1],
     [0,2],
     [1,2]] )


sess = tf.Session()
sess.run(tf.global_variables_initializer())

indices =  sess.run(B)

incre = tf.Variable(0)
template = tf.Variable(tf.zeros([6,2],tf.int32))

sess.run(tf.global_variables_initializer())

#There are 3 rows in the indices array
row = tf.gather( indices , [0,1,2])


for i in range(0, row.get_shape()[0] ) :

    newrow = tf.gather(row, i)

    exprow1 = tf.concat([tf.constant([i]), newrow[1:]], axis=0)
    exprow2 = tf.concat([tf.constant([i]), newrow[:1]], axis=0)

    template = tf.scatter_update(template, incre, exprow1)
    template = tf.scatter_update(template, incre + 1, exprow2)

    #Dataflow execution dependency is enforced.
    with tf.control_dependencies([template]): 
        incre = tf.assign(incre,incre + 2)

print(sess.run(tf.gather_nd(A,template)))

輸出是這個。

[11 12 15 13 18 17]