Was ist eine gute Methode, um eine Scheibe zu erhalten, die in TensorFlow teilweise von einem anderen Tensor indiziert wird?

Question

Angenommen, wir haben einen Tensor x mit unbekannter erster Dimension (zum Beispiel [?, 32, 32, 3]) und einen anderen Tensor i, der eigentlich ein Skalar ist. Gibt es eine gute Möglichkeit, das i -te Stück x nach erster Dimension zu erhalten, um beispielsweise einen Tensor der Dimension [32, 32, 3] zu erhalten? Ich bin neu bei TensorFlow und konnte mir nur diese äußerst plumpe Lösung einfallen lassen.

index = tf.concat(0, [i, tf.constant([0, 0, 0], tf.int64)]) 
size = [1, x.get_shape()[1].value, x.get_shape()[2].value, x.get_shape()[3].value] 
result = tf.unpack(tf.slice(x, index, size))[0] 

Answer 1

Sie können die Vorteile aus der Tatsache ziehen, dass -1 zum tf.slice()size Argumente ein besonderes Argument ist, „alle übrigen Elemente in dieser Dimension“ bedeutet. Dann wird ein Skalar i unter der Annahme (und nicht eine Länge-1-Vektor, wie es in Ihrem Code-Snippet zu sein scheint), können Sie tun:

result = tf.squeeze(tf.slice(x, tf.pack([index, 0, 0, 0]), [1, -1, -1, -1]), [0]) 

Alternativ können Sie tf.gather() verwenden eine oder mehrere Scheiben zur Auswahl ein Tensor in der nullten Dimension. In diesem Fall i muss ein Vektor sein:

i = tf.expand_dims(i, 0) # Converts `i` to a vector if it is a scalar. 
result = tf.squeeze(tf.gather(x, i), [0]) 

In beiden Fällen ist die tf.squeeze() op entfernt die 0-ten Dimension eines dreidimensionalen Ergebnis.