[python] Tensorflow tf.math.argmax 함수

Tensorflow에서 제공하는 tf.math.argmax함수는 tensor의 가장 큰 값의 index를 return해주는 함수이다.

 

기본 형태는

tf.math.argmax(input, axis = None, output_type = tf.dtype.int64, name = None)

위와 같다.

 

 input값은 tensor여야 하고, axis는 default는 0이고, max값을 찾을 축을 의미한다.

 

예제 1)

a = tf.constant([2, 20, 30, 6, 3])
tf.math.argmax(a)

a에서는 30이 가장 큰 값이고, 30의 index는 2 이므로, argmax 함수의 return 값은 2가 된다.

 

예제 2)

a = tf.constant([[2, 20, 30, 3, 6], [3, 11, 16, 1, 8], [14, 45, 23, 5, 27]])
tf.math.argmax(a)

tf.math.argmax(a)는 tf.math.argmax(a, 0)과 같다(default = 0)

a.shape = (3, 5)이고 axis = 0은 5를 의미하므로,

각각의 array에서 같은 index별로 가장 큰 값을 가진 array의 index를 return한다.

[2, 3, 14], [20, 11, 45], [30, 16, 23], [3, 1, 5], [6, 8, 27] 중 가장 큰 값의 index를 각각 return한다. 

return값은 [2, 2, 0, 2, 2] 가 된다.

 

예제 3)

a = tf.constant([[2, 20, 30, 3, 6], [3, 11, 16, 1, 8], [14, 45, 23, 5, 27]])
tf.math.argmax(a, 1)

a.shape = (3, 5)이고 axis = 1은 3를 의미하므로 axis=1을 parameter로 주게 되면, 

각각의 array [2, 20, 30, 3, 6], [3, 11, 16, 1, 8], [14, 45, 23, 5, 27] 에서 가장 큰 값의 index를 각각 return한다.

return값은 [2, 2, 1] 이 된다.