Wie kann ich mit bestimmten Indizes in TensorFlow auf Elemente eines 3D-Tensors zugreifen?
Ich versuche, die Zeilen eines 3D-Tensors in einer bestimmten Reihenfolge von Indizes zu erhalten. Hier sind die Eingaben:
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]]]
Ich kämpfe mit tf.gather_nd()
. Irgendein Vorschlag? Ich kann sehen, dass es hier passiert, bin mir aber nicht sicher, wie ich es auf die gesamte Matrix anwenden soll, ohne for
loop oder zu verwendentf.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]]
"""
EDIT: Ich habe eine ähnliche Frage in Bezug auf Numpy gestellt. Eine clevere Indizierungsantwort löst zwar die Numpy-Version, aber es ist schwierig, sie auf Tensoren anzuwenden. Schauen Sie sich hier die akzeptierte Antwort an: Wie kann ich Elemente aus der 3D-Matrix mit bestimmten Indizes in numpy abrufen?
Antworten
Duh, das war dumm! Es gibt bereits eine sehr gute Funktion, die für mehrdimensionale Arrays im Tensorflow funktioniert. tf.gather()
Schauen Sie sich das batch_dims Argument für weitere Informationen.
>> 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]]]