pytorch的argmax函数 python中argmax

目录

• 一、基本介绍
• 二、代码实验
• 1、一维数组情况
• 1.1、axis=0
• 1.2、axis=1
• 2、二维数组情况
• 2.1、axis=0
• 2.2、axis=1
• 3、三维数组情况
• 3.1、axis=0
• 3.2、axis=1
• 3.3、axis=2
• 3.4、axis=-1
• 四、Reference

一、基本介绍

numpy中的argmax简而言之就是返回最大值的索引，当使用np.argmax(axis)，这里方向axis的指定往往让人不理解。
简而言之：这里axis可以让我们从行、列、或者是深度方向来看一个高维数组。

二、代码实验

1、一维数组情况

简单一维情况，np.argmax()直接返回最大值的索引，不指定axis可以认为是将数组拉平之后寻找最大值的索引

1.1、axis=0

当我们指定axis=0时，其实是在列中作比较，寻找最大的行的索引

当然对于这个一维情况没有什么影响。

1.2、axis=1

当我们指定axis=1的时候报错了，这是因为我们的a是一维数组，没有axis=1这个轴，可见当我们使用np.argmax()时axis的指定不能超过所需要排序的数组

2、二维数组情况

不指定axis就是相当于把二维数组拉平，直接选取最大值的索引

2.1、axis=0

指定axis=0就是比较列，返回行索引中的最大值

我们改写一个b中的元素，我们期望的结果是[2,2,1,2]

实际结果和我们期望相符合

2.2、axis=1

指定axis=0就是比较行，返回列索引中的最大值

3、三维数组情况

一个三维数组可以视作一张图片，它的三个维度分别为(high, width, channels) 分别表示图像的高、宽、通道数(深度)。常见的彩色图像都有三个通道，我们以常见的RGB图像为例构建一个数组。

直接使用np.argmax()，就是之间将三维数组拉平，寻找最大值的索引

3.1、axis=0

单独查看c的三个通道的数据，如图所示

对于三个通道取axis=0意味分别比较列返回行的最大值索引

我们期望的返回值应该是[[1,1,1,],[1,1,1],[1,1,1]]，实际的结果和我们的期望一致

3.2、axis=1

对于三个通道取axis=1意味分别比较行返回列的最大值索引

我们期望的结果是[[2, 2, 2],[2, 2, 2],[2, 2, 2],[2, 2, 2]],，实际的结果和我们的期望一致

3.3、axis=2

取axis=2意味着我们从图像的深度方向(通道方向)来进行比较，可以认为三个数组的叠在一起的，分别对应channel0，channel1，channel2而我们取最大值的索引就是返回对应pixel像素所在的通道索引。

c的channel2所有的像素值均大于其他两个channel所有返回值应该是[[2,2,2,],[2,2,2,],[2,2,2,],[2,2,2,]]，实际结果和我的期望一致

3.4、axis=-1

axis=-1即是反过来看轴，对于三维情况axis=-1和axis=2一致

其他：

对于二维情况axis=-1和anxis=1一致

对于一维情况axis=0和anxis=-1一致

四、Reference