argmax用法

argmax是返回最大值的索引，第二个参数是对应的维数，0就是按第一维元素(行)取值，即同列的每一行，指定第二个参数为1，则按第二维（列）的元素取值，同行的每一列

import tensorflow as tf
import numpy as np
arr1=np.array([1,3,2,5,7,0])
arr2=np.array([[1,2,3],[3,2,1],[4,7,2],[8,3,2]])
print('arr1=',arr1)
print('arr2=',arr2)

argmax_1=tf.argmax(arr1)
#指定第二个参数0，按第一维元素(行)取值，即同列的每一行
argmax_20=tf.argmax(arr2,0) 
#指定第二个参数为1，则按第二维（列）的元素取值，同行的每一列
argmax_21=tf.argmax(arr2,1)
#指定第二个参数为-1，则最后维的元素取值
argmax_22=tf.argmax(arr2,-1)
with tf.Session() as sess:
    print(argmax_1.eval())
    print(argmax_20.eval())
    print(argmax_21.eval())
    print(argmax_22.eval())