Tensorflow中 tf.argmax()的用法

tf.argmax()的用法

tf.argmax()的参数形式为：

tf.argmax(array / matrix，axis=0 / 1)

括号内第二个参数 axis=0 时表示按列查找 列表/矩阵 中的最大元素的索引，axis= 1时表示按行查找 列表/矩阵 中最大元素的索引

比如：

import tensorflow as tf

A = tf.Variable([2,6,4,8,9,1])  #array
B = tf.Variable([[3,2,7],
				[2,9,8],
				[6,4,1]])       #matrix

A_1 = tf.cast(tf.argmax(A,0),tf.int64)   #类型转化
B_1 = tf.cast(tf.argmax(B,0),tf.int64)
B_2 = tf.cast(tf.argmax(B,1),tf.int64)
				
init = tf.global_variables_initializer()

with tf.Session() as sess:
	sess.run(init)
	print(A_1.eval())
	print(B_1.eval())
	print(B_2.eval())

结果：

4       #列表中最大数的索引       
[2 1 1]   #矩阵中每一列最大数的索引
[2 1 0]   #矩阵中每一行最大数的索引