概述
返回数组中最大元素的索引
该函数的用法与np.argmax()类似
语法
tf.math.argmax(
input,
axis=None
)
参数
- input：数组
- axis：轴
- axis = 0：返回 每一列 最大元素的 索引
- axis = 1：返回 每一行 最大元素的 索引
例子
例子
一维数组
import tensorflow as tf
l=[1,2,3,4,1,2]
res=tf.math.argmax(l)
print(res)
print(type(res))
print(int(res))
执行结果：
tf.Tensor(3, shape=(), dtype=int64)
<class 'tensorflow.python.framework.ops.EagerTensor'>
3
例子2
二维数组
np.random.seed(17)
arr=np.random.randint(1,20,(3,4))
print(arr)
执行结果如下：
[[16 18 16 7]
[17 14 5 8]
[13 8 2 15]]
按照axis=0执行，按列取
print(tf.math.argmax(arr,axis=0))
执行结果如下：
tf.Tensor([1 0 0 2], shape=(4,), dtype=int64)
按照axis=1执行，按列取
print(tf.math.argmax(arr,axis=1))
执行结果如下：
tf.Tensor([1 0 3], shape=(3,), dtype=int64)