Блог

Мне нужно извлечь значения из тензора на основе тензора индексов.

Мой код выглядит следующим образом:

 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)
  

Это работает медленно. Есть ли более эффективный способ?

Вы можете использовать метод 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]])>