如何根据我从张量流中的另一个矩阵获得的最大和次要值以及索引来获取矩阵中每一行的值?

时间:2018-07-18 03:38:13

标签: python tensorflow

如何根据我从张量流中另一个矩阵获得的最大和次要值索引获取矩阵中每一行的值?例如。我的矩阵张量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索引。

1 个答案:

答案 0 :(得分:2)

因此对于这种特定情况,此代码返回值。

  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]