Tensorflow中 tf.argmax()方法的用法

最新推荐文章于 2025-07-14 14:00:27 发布

原创最新推荐文章于 2025-07-14 14:00:27 发布 · 4.4k 阅读

11 ·

CC 4.0 BY-SA版权

TensorFlow 专栏收录该内容

9 篇文章

订阅专栏

tf.argmax()方法用于计算张量tensor沿着某一维度的最大值的索引。

tf.math.argmax(
    input,
    axis=None,
    name=None,
    dimension=None,
    output_type=tf.int64)

第一个参数 input：输入的需要处理的张量 tensor，必须是以下数据类型之一float32, float64, int32, uint8, int16, int8, complex64, int64, qint8, quint8, qint32, bfloat16, uint16, complex128, half, uint32, uint64
第二个参数 axis：指定的维（或轴），取值为0，代表延列取最大值索引，取值为1，表示延行取最大值索引，如果不指定，则默认为按照列计算，取值为2,3,4...是，则为多维函数的相应维度，以此类推
第三个参数 name：操作的名称;
第四个参数 dimension：已弃用;
第五个参数 output_type：返回值的数据类型，是tf.int32, tf.int64之一。默认为tf.int64;

以一个维度是2，形状是[2,3]的tensor举例，结果：

import tensorflow as tf
 
x = [[1,2,3],
     [2,1,3]]
 
xx = tf.cast(x,tf.float32)
 
idx_default = tf.argmax(xx)
idx_0 = tf.argmax(xx, axis=0)
idx_1 = tf.argmax(xx, axis=1)
 
 
with tf.Session() as sess:
    i_d,i_0,i_1 = sess.run([idx_default, idx_0, idx_1])
 
print(i_d)
print(i_0)
print(i_1)
 
 
###  输出结果为 ###
# [1 0 0]
# [1 0 0]
# [2 2]

类似的方法还有：

tf.reduce_mean ：计算tensor指定轴方向上的所有元素的累加和;
tf.reduce_sum ：计算tensor指定轴方向上的所有元素的累加和;
tf.reduce_max  :  计算tensor指定轴方向上的各个元素的最大值;
tf.reduce_all :  计算tensor指定轴方向上的各个元素的逻辑和（and运算）;
tf.reduce_any:  计算tensor指定轴方向上的各个元素的逻辑或（or运算）;