TensorFlow에서 지정된 인덱스를 사용하여 3D 텐서의 요소에 액세스하려면 어떻게해야합니까?

Question

특정 인덱스 순서로 3D 텐서의 행을 가져 오려고합니다. 입력 내용은 다음과 같습니다.

import tensorflow as tf

matrix = tf.constant([
    [[0, 1], [2, 3], [4, 5], [6, 7]], 
    [[8, 9], [10, 11], [12, 13], [14, 15]], 
    [[16, 17], [18, 19], [20, 21], [22, 23]], 
    [[24, 25], [26, 27], [28, 29], [30, 31]], 
    [[32, 33], [34, 35], [36, 37], [38, 39]]
])

indx = tf.constant([[3,2,1,0], [0,1,2,3], [1,0,3,2], [0,3,1,2], [1,2,3,0]])

# required output tensor:
[[[6, 7], [4, 5], [2, 3], [0, 1]],
 [[8, 9], [10, 11], [12, 13], [14, 15]],
 [[18, 19], [16, 17], [22, 23], [20, 21]],
 [[24, 25], [30, 31], [26, 27], [28, 29]],
 [[34, 35], [36, 37], [38, 39], [32, 33]]]

나는 tf.gather_nd(). 어떠한 제안? 여기에서 일어나는 것을 볼 수 있지만 루프 를 사용하지 않고 전체 매트릭스에 적용하는 방법을 잘 모르겠습니다.fortf.map_fn

print(tf.gather_nd(matrix[0], tf.expand_dims(indx, -1)[0]).numpy().tolist())
print(tf.gather_nd(matrix[1], tf.expand_dims(indx, -1)[1]).numpy().tolist())
print(tf.gather_nd(matrix[2], tf.expand_dims(indx, -1)[2]).numpy().tolist())
print(tf.gather_nd(matrix[3], tf.expand_dims(indx, -1)[3]).numpy().tolist())
print(tf.gather_nd(matrix[4], tf.expand_dims(indx, -1)[4]).numpy().tolist())

"""
[[6, 7], [4, 5], [2, 3], [0, 1]]
[[8, 9], [10, 11], [12, 13], [14, 15]]
[[18, 19], [16, 17], [22, 23], [20, 21]]
[[24, 25], [30, 31], [26, 27], [28, 29]]
[[34, 35], [36, 37], [38, 39], [32, 33]]
"""

편집 : 나는 numpy와 관련하여 비슷한 질문을했습니다. 영리한 인덱싱 답변은 numpy 버전을 해결하지만 Tensor에 적용하기는 어렵습니다. 여기에서 허용되는 답변을 자유롭게 살펴보십시오 : numpy에서 지정된 인덱스를 사용하여 3D 매트릭스에서 요소를 어떻게 얻을 수 있습니까?

Answer 1

멍청 했어! tensorflow의 다차원 배열에서 작동하는 매우 훌륭한 함수가 이미 있습니다. 자세한 내용 tf.gather()은 batch_dims 인수를 확인하십시오.

>> tf.gather(matrix, indx, batch_dims=1).numpy().tolist()
[[[6, 7], [4, 5], [2, 3], [0, 1]],
 [[8, 9], [10, 11], [12, 13], [14, 15]],
 [[18, 19], [16, 17], [22, 23], [20, 21]],
 [[24, 25], [30, 31], [26, 27], [28, 29]],
 [[34, 35], [36, 37], [38, 39], [32, 33]]]