官方API的定义
tf.argmax(input, axis=None, name=None, dimension=None)
Returns the index with the largest value across axes of a tensor.
Args:
input: A Tensor. Must be one of the following types: float32, float64, int64, int32, uint8, uint16, int16, int8, complex64, complex128, qint8, quint8, qint32, half.
axis: A Tensor. Must be one of the following types: int32, int64. int32, 0 <= axis < rank(input). Describes which axis of the input Tensor to reduce across. For vectors, use axis = 0.
name: A name for the operation (optional).
Returns:
A Tensor of type int64.
关于axis的说明
input可以是数组,矩阵,当input是矩阵(np.mat)时:
- axis = 0的时候返回每一列最大值的位置索引
- axis = 1的时候返回每一行最大值的位置索引
import tensorflow as tf
import numpy as np
x = np.mat([0, 0, 1])
y = np.mat([[0, 0, 1], [1, 2, 3]])
x_max_index = tf.argmax(x, 1)
y_max_index = tf.argmax(y, 1)
with tf.Session() as sess:
x_index = sess.run(x_max_index)
y_index = sess.run(y_max_index)
print("x_max_index = ", x_index)
print("y_max_index = ", y_index)
输出:
x_max_index = [2]
y_max_index = [2 2]
输入和输出的数据类型一样,当输入是array时,输出也是array。不一样的地方在于,当数组是一维数组时,axis只能为0,输出最大数的索引值。