用 numpy 的人十个里有八个在第一次面对 argmax 的 axis 参数时都栽过跟头。我在实际项目和带新人的过程中见过太多人把 axis0 当成“取每一行的最大值”也见过有人在三维数组面前直接懵掉还有人压根不传 axis结果返回的索引和预期差出十万八千里。这篇文章就把 numpy.argmax 的 axis0、axis1、axis-1 彻底掰开揉碎讲清楚每一个取值在这个函数里到底是怎么作用的。适合刚接触 numpy 的初学者也适合用过一阵子、但始终没搞懂 axis 本质的进阶用户。看完之后你不需要死记任何口诀一条“形状消去”的规则就能推导出所有情况。1. argmax 到底在干什么找位置不是找值1.1 不传 axis 时的默认行为一个最简单的例子import numpy as np a np.array([3, 1, 4, 1, 5]) print(np.argmax(a)) # 4 print(np.max(a)) # 5数组 a 里最大的元素是 5它出现在下标 4 的位置所以 np.argmax(a) 返回 4。注意这里有个很容易混淆的点argmax 返回的是“位置”而 max 返回的是“数值”。很多人会下意识以为 argmax 也会返回 5这就是第一个思维误区。这个“返回位置”的设计在工程里非常关键。做分类任务的时候模型的输出经常是一个概率向量比如 [0.1, 0.7, 0.2]我们真正想要的不是“最大概率是 0.7”而是“概率最大的那个类别是第几个”。max 告诉你有多大概率argmax 告诉你去哪儿找这个概率。凡是需要“定位”的功能比如找峰值位置、找最佳匹配项、找预测类别都离不开 argmax。不传 axis 的时候numpy 会先把数组按 C 顺序展开成一维再返回全局最大值的扁平下标。这个行为在二维及以上的数组上尤其危险后面我会专门讲这个坑。1.2 传了 axis 之后输出形状怎么变axis 参数一旦指定语义就完全不同了。这时候 argmax 不再全局找最大值而是沿着你指定的那条轴每次滑动一步取一个“局部最大值位置”。如果只看结果形状有一条非常硬核的规则指定 axisk 之后返回结果的形状 原数组形状去掉第 k 维。举例来说一个二维数组的形状是 (2, 3)——2 行 3 列。axis0 时结果形状是 (3,)axis1 时结果形状是 (2,)axis-1 在二维里等价于 axis1结果形状同样是 (2,)。你看无论选哪个轴结果都比原数组少一维这就是“压缩”的含义。我把这条规则叫做“形状消去法”。它是所有 axis 题目的总纲比任何口诀都可靠。后面遇到再高维的数组只要先想清楚“去掉这一维之后剩下什么形状”你就不会把 axis 的方向搞反。1.3 axis 的本质沿着哪条轴做比较很多人记 axis 靠背“axis0 是行axis1 是列”这句话本身没错但到了 argmax 里就特别容易产生误导。“axis0 是行”指的是第 0 维是行方向但如果要你“沿着第 0 维做比较”实际动作其实是固定每一列在这一列内部从上到下遍历把这一列压成一个位置号。也就是说axis0 的结果是“每一列的最大值所在行号”。axis1 同理固定每一行在这一行内部从左到右遍历结果是“每一行的最大值所在列号”。看明白这个逻辑了吗axis 参数指定的是“你要消掉哪一个维度”。二维数组里轴 0 是行方向消掉行方向剩下的每一列出一个结果轴 1 是列方向消掉列方向剩下的每一行出一个结果。这个理解方式能直接推广到任意维度而“背行、背列”的方式到了三维就彻底崩盘。2. 二维矩阵手算演示axis0 和 axis1 到底谁是谁2.1 一个能口算的矩阵纸上谈兵没有用咱们拿一个真实矩阵手算一遍。我特意选了个小矩阵每个值都能心算arr np.array([[8, 1, 6], [3, 5, 7], [4, 9, 2]])这个数组形状是 (3, 3)。先算 axis0规则是“沿着第 0 维行方向比较消掉行维度结果为每一列的最大值行号”。逐列看第 0 列是 [8, 3, 4]最大值 8 在第 0 行所以结果是 0第 1 列是 [1, 5, 9]最大值 9 在第 2 行所以结果是 2第 2 列是 [6, 7, 2]最大值 7 在第 1 行所以结果是 1。最终得到 [0, 2, 1]。用代码验证print(np.argmax(arr, axis0)) # [0 2 1]注意结果里每个数字的含义它表示“这一列的最大值出现在第几行”。这个数字是行号但结果本身是“按列排列”的。2.2 再手算 axis1axis1 的规则是“沿着第 1 维列方向比较消掉列维度结果为每一行的最大值列号”。逐行看第 0 行是 [8, 1, 6]最大值 8 在第 0 列所以结果是 0第 1 行是 [3, 5, 7]最大值 7 在第 2 列所以结果是 2第 2 行是 [4, 9, 2]最大值 9 在第 1 列所以结果是 1。结果是 [0, 2, 1]print(np.argmax(arr, axis1)) # [0 2 1]这个例子挺有意思两个轴算出来的结果碰巧一样。我故意选了这个矩阵也提醒你实战里两个轴的结果往往完全不同别指望它们偶然相等就能蒙混过关。2.3 一个不靠背、靠“想象”的记忆法我给学生讲 axis 的时候偏爱一个比喻。你可以把一个二维数组想象成一片矩形场地每个位置上立着一根柱子柱子高度就是数组里的数值。axis0想象有一束光从正上方竖直打下来。光穿过每一列在这一列里找到最高的那根柱子记录它的行号。从上往下看每一列都变成一个点原来 (3,3) 的场地就缩成了 3 个点对应形状 (3,)。axis1想象光从侧面水平照过来。光扫过每一行在这一行里找到最高的柱子记录它的列号。每一行缩成一个点结果形状是 (3,)。这个“顶灯、侧灯”的类比能帮你直接绕过“行还是列”的口诀纠缠。再用“形状消去法”兜底axis0 消掉行axis1 消掉列。两个方法互相印证基本就不会再混了。3. axis-1 到底是什么三维数组又该怎么看3.1 三维数组的三个轴到了三维光靠“行、列”已经说不清了。一个形状为 (2, 3, 4) 的数组可以理解成“2 个 3 行 4 列的矩阵叠在一起”。三个轴分别是axis0深度方向也就是“第几层矩阵”结果形状 (3, 4)axis1行方向也就是“矩阵里的第几行”结果形状 (2, 4)axis2列方向也就是“矩阵里的第几列”结果形状 (2, 3)。用前面那套“形状消去法”就是在 2、3、4 这三个数里划掉一个。axis0 划掉最外面的 2剩下 (3,4)axis1 划掉 3剩下 (2,4)axis2 划掉 4剩下 (2,3)。划谁消谁一目了然。3.2 用递增数组验证你的理解口说无凭我造一个非常“心算友好”的三维数组。让每个矩阵里的数字都按行、按列递增arr3d np.array([ [[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12]], [[13, 14, 15, 16], [17, 18, 19, 20], [21, 22, 23, 24]], ]) print(arr3d.shape) # (2, 3, 4)这个数组有个特点每个维度的值都单调递增所以最大值的位置完全可以口算。axis0比较两个矩阵的对应位置第二个矩阵的每一个数字都比第一个大所以每个位置的最大值都在第 1 层结果是一个全 1 的 (3,4) 矩阵axis1比较每个矩阵内部的三行第三行最大所以每个位置的最大值都在行号 2结果是全 2 的 (2,4) 矩阵axis2比较每一行内部的四个数字第四个最大结果是全 3 的 (2,3) 矩阵。跑一下什么都明白了print(np.argmax(arr3d, axis0)) # [[1 1 1 1] # [1 1 1 1] # [1 1 1 1]] print(np.argmax(arr3d, axis1)) # [[2 2 2 2] # [2 2 2 2]] print(np.argmax(arr3d, axis2)) # [[3 3 3 3] # [3 3 3 3]]拿这种“全 1、全 2、全 3”的结果做自我检查比自己随便造一个数组高效得多。你要是记错了 axis一眼就能看出来。3.3 axis-1 指向的是最后一根轴axis-1 一点也不玄乎它跟 Python 里的负索引是同一个逻辑。arr.shape[-1] 表示最后一维的长度那么 axis-1 就表示“沿着最后一维做比较”。放到具体维度上一维数组axis-1 等价于 axis0二维数组axis-1 等价于 axis1三维数组axis-1 等价于 axis2四维数组axis-1 等价于 axis3。所以 axis-1 不是一个独立的新模式它只是“懒人写法”。我推荐大家写代码统一用 axis-1 而不是 axis1 或者 axis2原因很简单很多时候你的数据到底有没有 batch 维、中间有没有插进来别的维度是不确定的。写 -1 意味着“永远取最后一根轴作为特征维度”代码的坚韧度会高很多业务里加了一维也不用回头改。4. 实际场景里的标准写法分类、通道、序列4.1 多分类预测argmax(axis-1) 取类别最典型的使用场景就是分类任务。模型一般会输出一个形状为 (batch, num_classes) 的概率矩阵每一行是一个样本每一列是一个类别。这时候要得到每个样本的预测类别写 argmax(axis-1) 就好# 模拟 5 个样本、3 个类别的概率输出 probs np.array([ [0.1, 0.7, 0.2], [0.8, 0.1, 0.1], [0.2, 0.2, 0.6], [0.4, 0.5, 0.1], [0.3, 0.3, 0.4], ]) preds np.argmax(probs, axis-1) print(preds) # [1 0 2 1 2]这里每个数字就是“该样本概率最大的类别下标”。如果类别列表是 [猫, 狗, 鸟]那么第一个样本预测结果是下标 1也就是“狗”。把概率矩阵想象成场地上的一排排柱子axis-1 就是沿着每一行找最高的柱子然后把行号记录下来。有的同学会问这里写 axis1 不也行吗没错在二维情况下两者完全等价。但如果你某天把数据改成了 (seq_len, batch, num_classes) 之类的结构写 axis1 就错到姥姥家了写 axis-1 依然正确。这就是“最后一根轴”写法的价值。4.2 图像通道和序列特征超过二维也别慌图像数据常见的布局是 (batch, height, width, channels)也就是“一批图每张图有高、宽、通道三个维度”。如果要做逐像素的分类比如语义分割就得对每个像素在通道维度上取概率最大的那个通道下标。通道是最后一维所以依然是 argmax(axis-1)# 模拟 2 张图每张 3x34 个通道 seg_map np.random.rand(2, 3, 3, 4) pred_pixel np.argmax(seg_map, axis-1) print(pred_pixel.shape) # (2, 3, 3)结果是 (batch, height, width)每个像素位置存放的是通道下标。你看用“形状消去法”一划axis-1 把 4 划掉剩下 (2,3,3)正好就是“每个像素一个类别号”的形状。时间序列数据也一样。假设数据形状是 (batch, time_steps, features)每个时间步有很多个特征维度想找出每个时间步上最显著的特征就是 argmax(axis-1)结果形状变成 (batch, time_steps)。这类写法在序列模型的特征分析里非常常见。4.3 和 max、where 的区别三个函数别混用argmax 和 max 最容易混淆。np.max 返回的是数值本身np.argmax 返回的是位置。两者经常搭配使用arr np.array([3, 1, 4, 1, 5]) print(np.max(arr)) # 5 print(np.argmax(arr)) # 4还有一个 np.where(arr np.max(arr)) 也能拿到最大值的位置但它返回的是所有满足条件的位置结果是一个元组每个元素对应一个轴上的下标数组而不是单个整数。而且它找的是“所有最大值”包括重复的argmax 则只认第一个。比如数组 [1, 5, 3, 5]argmax 返回 1而 where 会返回两个位置。另外有个隐藏坑默认情况下 numpy 的 argmax 在遇到多个最大值时返回的是“第一次出现”的位置。这个规则在不同语言、不同框架里不一定一致所以如果你在做跨语言移植务必把这个行为查清楚不然排查半天都找不到根因。4.4 新版本里可选的 keepdims 参数较新版本的 numpy 给 argmax 加了一个 keepdims 参数设为 True 时被消掉的那个维度会以长度为 1 的形式保留下来。举个例子arr np.array([[8, 1, 6], [3, 5, 7], [4, 9, 2]]) print(np.argmax(arr, axis0, keepdimsTrue)) # [[0 2 1]] print(np.argmax(arr, axis0, keepdimsTrue).shape) # (1, 3)结果从 (3,) 变成了 (1, 3)。这个特性在做广播运算时特别有用能让结果直接和原数组或者其他同形状数组对齐。不过日常代码里用不带 keepdims 的写法更多这里提一句免得你看到别人的代码里多出这个参数时一头雾水。5. 常见坑与自查技巧实录5.1 最经典的坑忘了传 axis整个数组被展平二维数组不传 axis 时numpy 会把整个数组按 C 顺序展平成一个一维数组再找全局最大值。这个“全局位置”往往不是你想要的“每行最大值位置”或者“每列最大值位置”。看个例子arr np.array([[8, 1, 6], [3, 5, 7], [4, 9, 2]]) print(np.argmax(arr)) # 7数组展平后是 [8, 1, 6, 3, 5, 7, 4, 9, 2]最大值 9 的下标是 7对应的原始位置是第 2 行第 1 列。如果你原本想找每行的最大值这个结果毫无意义。所以只要你的数据不是真·一维就一定把 axis 写清楚哪怕你觉得“反正不写也能跑”。5.2 第二坑把 axis0 理解成“每行最大值”axis0 确实是“行方向”但它作用的方向和你以为的方向是反的。axis0 消掉的是行维度所以结果对应的是“每一列”而不是“每一行”。我把这个逻辑再强调一次axis0一列一列地看每列返回一个行号axis1一行一行地看每行返回一个列号。很多人在这两个方向上摔跤根源是把“方向”和“结果归属”混为一谈。axis0 表示沿行方向走但走到头消掉的是行维留下来的是每列。用“形状消去法”再验证一遍二维数组 (2,3)axis0 划掉 2剩下 33 正好是列数结果当然是一列一个值。想明白这一点方向就再也不会反了。5.3 第三坑三维以上只知道 -1不知道它到底消了哪根轴axis-1 是好习惯但如果你完全不知道 -1 指向哪儿遇到数据形状变化一样会踩坑。比如你的数据最后一位恰好不是特征维而是一个 id 维或者时间维那 argmax(axis-1) 的结果就会完全偏离预期。所以我建议写 -1 之前先打印一下 arr.shape看清楚最后一维的长度和含义。使用 -1 是建立在“我知道最后一维是什么”的前提下不是无脑偷懒。5.4 自查法每次写完先对形状再手算一个值我自己的习惯非常固定写带 axis 的 argmax 时第一件事就是看结果形状。规则很简单原数组去掉第 k 维剩下的就是输出形状。比如输入 (4, 6, 5)axis1输出只能是 (4, 5)。形状不对后面全白搭越早检查越省钱。形状对上了再随便挑一个位置手算验证。比如二维数组取 axis1我可以手动固定第 2 行数一数这一行的最大值在第几列再对齐输出数组里第 2 个元素的值。两次对上了这个调用才算真正安全。5.5 需要完整坐标时怎么把消掉的维度找回来argmax(axisk) 返回的只是“被消掉的第 k 维上的坐标”其他维度的坐标是隐含的。比如 (3,3) 数组取 axis1结果 [0,2,1] 只告诉了你每行最大值的列号没告诉你是哪一行。想要完整的 (row, col) 坐标需要自己补上行号rows np.arange(arr.shape[0]) cols np.argmax(arr, axis1) coords np.column_stack((rows, cols)) print(coords) # [[0 0] # [1 2] # [2 1]]如果一开始就不指定 axis直接 np.argmax(arr) 拿到的是展平后的下标这时候可以用 np.unravel_index 还原成多维坐标。这两种做法各有适用场景但核心就一句话argmax 永远只告诉你“沿轴方向的位置”别的坐标要自己算。最后我给个小建议以后凡是看到别人代码里出现 np.argmax(x, axis-1)你不妨先看一眼 x.shape再在心里过一遍“划掉最后一维后剩什么形状”两秒就能判断这个调用合不合理。这个习惯我坚持了很多年帮我躲过了大量低级但极其隐蔽的索引错误。argmax 本身不难难的是把 axis 的“消维”本质刻进脑子。不需要死记任何口诀你只需要记住一条规则axis 指向哪一根轴哪一根轴就会被消掉返回的是剩下每一份数据里最大值在那根轴上的坐标。
