Как эффективно получить значения каждой строки в тензоре с помощью индексов?

Question

У меня есть тензор my_tensor с формой [batch_size, seq_length], и у меня есть другой тензор с именем idx с формой [batch_size, 1], который состоит из индексов, которые начинаются с 0 и заканчиваются на "seq_length".

Я хочу извлечь значения в каждой строке my_tensor, используя индексы, определенные в idx.

Я пытался использовать tf.gather_nd и tf.gather, но мне это не удалось.

Рассмотрим следующий пример:

batch_size = 3
seq_length = 5
idx = [2, 0, 4]

my_tensor = tf.random.uniform(shape=(batch_size, seq_length))

Я хочу получить значения в

[[0, 2],
 [1, 0],
 [3, 4]]

от my_tensor.

Я должен выполнить дальнейшую обработку над ними, поэтомуЯ хотел бы иметь их в то же время (я не знаю, возможно ли это вообще) и эффективным способом;однако я не смог придумать другие методы.

Я ценю любую помощь:)

Answer 1

Хитрость заключается в том, чтобы сначала преобразовать ваш набор индексов в логическую маску, которую затем можно использовать для уменьшения my_tensor, как вы описали, используя операцию boolean_mask .

Вы можетеВыполните это с помощью однократного кодирования idx тензора.

Итак, где idx = [2, 0, 4] мы можем сделать tf.one_hot(idx, seq_length), чтобы преобразовать его во что-то вроде этого:

[ [0., 0., 1., 0., 0.],
  [1., 0., 0., 0., 0.],
  [0., 0., 0., 0., 1.] ]

Затем, сложив все вместе, скажем my_tensor:

[ [0.6413697 , 0.4079175 , 0.42499018, 0.3037368 , 0.8580252 ],
  [0.8698617 , 0.29096508, 0.11531639, 0.25421357, 0.5844104 ],
  [0.6442119 , 0.31816053, 0.6245482 , 0.7249261 , 0.7595779 ] ]

мы можем действовать следующим образом:

result = tf.boolean_mask(my_tensor, tf.one_hot(idx,seq_length))

, чтобы получить:

[0.42499018, 0.8698617 , 0.7595779 ]

, как ожидается