How to get the indices of the maximum value in a tensor using Tensorflow?

user_13919166 0
Viewed 1467

I am looking for the best and optimized way (without loops) to get the indices of the maximum value of tensor (Rank 2 tensor) using TensorFlow 1.14, in Google Colab i tried this code using TF2.0:

import tensorflow as tf
A = tf.constant([[0.2,0.8],[0.3,0.9],[0.4,0.7],[0.5,0.4]])
b = tf.math.argmax(A,0)
bb = b.numpy()

The indices of the max here is [1,1], but the problem is i have to give the axis as input, and it doesn't give me the right one even i change the axis.

tensorflow numpy python
3 Answers

You could use Tensorflow functions to do that as:

max_val = tf.reduce_max(A, keepdims=True)
cond = tf.equal(A, max_val)
res = tf.where(cond)
res
# <tf.Tensor: shape=(1, 2), dtype=int64, numpy=array([[1, 1]], dtype=int64)>

If you want a 1D array in result, add the following:

res_1d = tf.squeeze(res)
res_1d
# <tf.Tensor: shape=(2,), dtype=int64, numpy=array([1, 1], dtype=int64)>

I have not used tf 1.14, but I guess you cannot use .numpy() for A[res_1d.numpy()[0]]. But you could do the following:

tf.slice(A, res_1d, [1, 1])
# <tf.Tensor: shape=(1, 1), dtype=float32, numpy=array([[0.9]], dtype=float32)>
user_137205650

Don't know if it's the best way to do it but I've found this (from the numpy documentation) that might help you :

import numpy as np

c = A.numpy()
np.unravel_index(np.argmax(c, axis=None), c.shape) # outputs (1,1)
user_137069040

Tou could reshape your nd-array into a 1d-array and do argmax. Then calculate the true index in your nd-array

user_18331750
Related