如何根据张量流中的索引张量填充值张量？

Question

我需要根据索引张量从张量中提取值。我的代码如下：arr = tf.constant([10, 11, 12]) # array of valuesinds = tf.constant([0, 1, 2])   # indicesres = tf.map_fn(fn=lambda t: arr[t], elems=inds)它运作缓慢。有更有效的方法吗？

Answer 1

您可以使用 tf.gather 方法

  arr = tf.constant([10, 11, 12]) # array of values

  inds = tf.constant([0, 2]) 

  r = tf.gather(arr , inds)#<tf.Tensor: shape=(2,), dtype=int32, numpy=array([10, 12])>

如果您有一个多维张量，则 tf.gather 有一个“axis”参数来指定检查索引的维度：

arr = tf.constant([[10, 11, 12] ,[1, 2, 3]]) # shape(2,3)

inds = tf.constant([0, 1]) 

# axis == 1

r = tf.gather(arr , inds , axis = 1)#<tf.Tensor: shape=(2, 2), dtype=int32, numpy=array([[10, 11],[ 1,  2]])>

 

# axis == 0

 r = tf.gather(arr , inds , axis = 0) #<tf.Tensor: shape=(2, 3), dtype=int32, numpy=array([[10, 11, 12], [ 1,  2,  3]])>