Tensorflow：使用argmax切片张量

Question

I have a tensor with shape tf.shape(t1) = [1, 1000, 400] and I obtain the indices of the maxima on the 3rd dimension using max_ind = tf.argmax(t1, axis=-1) which has shape [1, 1000] . Now I have a second tensor that has the same shape as t1 : tf.shape(t2) = [1, 1000, 400] .

I want to use the maxima indices from t1 to slice t2 so the output has the form

[1, 1000]

A more visual description: The resulting tensor should be like the result of tf.reduce_max(t2, axis=-1) but with the location of the maxima in t1

Answer 1

You can achieve this through tf.gather_nd , although it is not really straightforward. For example,

shape = t1.shape.as_list()
xy_ind = np.stack(np.mgrid[:shape[0], :shape[1]], axis=-1)
gather_ind = tf.concat([xy_ind, max_ind[..., None]], axis=-1)
sliced_t2 = tf.gather_nd(t2, gather_ind)

If on the other hand the shape of your input is unknown as graph construction time, you could use

shape = tf.shape(t1)
xy_ind = tf.stack(tf.meshgrid(tf.range(shape[0]), tf.range(shape[1]),
                              indexing='ij'), axis=-1)

and the remainder is the same as above.