Get indices of maximum values in Tensorflow

Issue

I would like to get the indices of maximum values.

Eg:

[
 [
  [0.1 0.3 0.6],
  [0.0 0.4 0.1]
               ],
 [
  [0.9 0.2 0.6],
  [0.8 0.1 0.5]
               ]
                ]

I would like to get [[0,0,2], [0,1,1], [1,0,0], [1,1,0]]. How do I do that in the easiest way in Tensorflow?

Solution

You can take advantage of TF’s broadcast in the last dimension

a = tf.constant([[[0.1, 0.3, 0.6],[0.0, 0.4, 0.1]],[[0.9, 0.2, 0.6],[0.8, 0.1, 0.5]]])
b = tf.reduce_max(a, -1, keepdims=True)
tf.where(a == b)

Output

<tf.Tensor: shape=(4, 3), dtype=int64, numpy=
array([[0, 0, 2],
       [0, 1, 1],
       [1, 0, 0],
       [1, 1, 0]], dtype=int64)>

In case of multiple max values per row and you only want to keep index of the first, you can derive which segment each row in the result corresponds to, then do a segment_min to get the first index in each segment.

a = tf.constant([[[0.1, 0.6, 0.6],[0.0, 0.4, 0.1]],[[0.9, 0.2, 0.6],[0.8, 0.1, 0.5]]])
b = tf.reduce_max(a, -1, keepdims=True)
c = tf.cast(tf.where(a == b), tf.int32)
d = tf.reduce_sum(tf.math.cumprod(a.shape[:-1], reverse=True, exclusive=True) * c[:,:-1], axis=1)
tf.math.segment_min(c,d)

Output

<tf.Tensor: shape=(4, 3), dtype=int32, numpy=
array([[0, 0, 1],
       [0, 1, 1],
       [1, 0, 0],
       [1, 1, 0]])>

Answered By – bui

This Answer collected from stackoverflow, is licensed under cc by-sa 2.5 , cc by-sa 3.0 and cc by-sa 4.0

Leave a Reply

(*) Required, Your email will not be published