eleganter Weg,

Question

ein Element pro Zeile in Tensorflow auswählen

Gegeben ...

• eine Matrix A der Form [m, n]
• ein Tensor I der Form [m]

ich eine Liste erhalten möchten J von Elementen aus A, wobei J[i] = A[i, I[i]].

Das heißt, I enthält den Index des Elements zur Auswahl aus jeder Zeile in A.

Kontext: Ich habe bereits die argmax(A, 1) und jetzt möchte ich auch die max. Ich weiß, dass ich einfach reduce_max verwenden kann. Und nach ein wenig versucht, um kam ich auch mit diesem nach oben:

J = tf.gather_nd(A, 
    tf.transpose(tf.pack([tf.to_int64(tf.range(A.get_shape()[0])), I]))) 

Wo die to_int64 benötigt wird, weil Bereich produziert nur int32 und argmax produziert nur int64.

Keines der beiden kommt mir besonders elegant vor. Man hat Laufzeit Overhead (wahrscheinlich über den Faktor n) und der andere hat einen unbekannten Faktor kognitiven Overhead. Fehle ich hier etwas?

Answer 1

Link zur Verfügung gestellt von @ yaroslav-bulatov mentiones this Lösung:

def get_elements(data, indices): 
    indeces = tf.range(0, tf.shape(indices)[0])*data.shape[1] + indices 
    return tf.gather(tf.reshape(data, [-1]), indeces) 

Ihre Lösung ist nicht differenzierbar (weil Gradienten für tf.gather_nd werden derzeit nicht unterstützt).

Hoffentlich wird data[:, indices] bald vorgestellt werden.

Answer 2

Dies ist eine eher späte Antwort, aber konnte

mask = tf.one_hot(I, depth=n, dtype=tf.bool, on_value=True, off_value=False) 
elements = tf.boolean_mask(A, mask) 

Erledigen zu tun, was Sie suchen?

edit: Ich sollte darauf hinweisen, dass dies keine gute Idee ist, wenn A bereits ein sehr großer Tensor ist, da dies zu einer dichten Matrix führt.