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] #矩阵中每一行最大数的索引