python - 如何根据 tensorflow 中的索引张量填充值张量？

Question

我需要根据索引张量从张量中提取值。
我的代码如下:

arr = tf.constant([10, 11, 12]) # array of values
inds = tf.constant([0, 1, 2])   # indices
res = 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 有一个“轴”参数来指定您检查索引的维度:

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]])>