tf.argmax()用法

tf.argmax(vector, 1)：返回的是vector中的最大值的索引號，如果vector是一個向量，那就返回一個值，如果是一個矩陣，那就返回一個向量，這個向量的每一個維度都是相對應矩陣行的最大值元素的索引號。

import tensorflow as tf  
import numpy as np  

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

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

輸出：

[4]
[2 1]