1. 程式人生 > >tf.argmax(vector, 1)

tf.argmax(vector, 1)

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

import tensorflow as tf
A = tf.constant([[1,2,3,4,5]])
B = tf.constant([[1,2,3],[6,5,4]])
with tf.Session() as sess:
    print(sess.run(tf.argmax(A,1)))
    print(sess.run(tf.argmax(
B,1)))

輸出
[4]
[2 0]