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

如何基于张量流中的索引张量填充值的张量？

1 个答案: