TensorFlow tf.argmax()

tf.argmax(vector, dimension)：返回的是vector中的最大值的索引号，如果vector是一个向量，那就返回一个值，如果是一个矩阵，那就返回一个向量，这个向量的每一个维度都是相对应矩阵行的最大值元素的索引号。

简单说：
dimension=0 按列找
dimension=1 按行找
tf.argmax()返回最大数值的下标

举例：

import tensorflow as tf

A = [[1, 3, 4, 5, 6]]
B = [[1, 3, 4], [2, 4, 1]]

with tf.Session() as sess:
    print(sess.run(tf.argmax(A, 0)))
    print(sess.run(tf.argmax(A, 1)))
    print(sess.run(tf.argmax(B, 0)))
    print(sess.run(tf.argmax(B, 1)))

打印：

[0 0 0 0 0]
[4]
[1 1 0]
[2 1]