短小精悍算例：TensorFlow中tf.argmax()函数的使用-返回最大值的位置索引

import numpy as np 
import tensorflow as tf
x = np.array([[3, 1, 2],
              [4, 7, 3],
              [5, 0, 1], 
              [2, 4, 6]])
a = tf.argmax(x, axis=0)  # 求各列最大值
b = tf.argmax(x, axis=1)  # 求各行最大值

sess = tf.Session()
print(sess.run(a))
print(sess.run(b))

输出结果:
[2 1 3]
[0 1 0 2]