argmax最经典解释

2. 高维
以三维为例，计算思路与二维相同。

三维计算之后降维，将返回一个二维数组。

一个m×n×p维的矩阵，
axis为0，舍去m，返回一个 n×p 维的矩阵
axis为1，舍去n，返回一个 m×p 维的矩阵
axis为2，舍去p，返回一个 m×n 维的矩阵

three_dim_array = [[[1, 2, 3, 4],  [-1, 0, 3, 5]],
                   [[2, 7, -1, 3], [0, 3, 12, 4]],
                   [[5, 1, 0, 19], [4, 2, -2, 13]]]
a = np.argmax(three_dim_array, axis = 0)
print(a)
b = np.argmax(three_dim_array, axis = 1)
print(b)
c = np.argmax(test, axis = 2)
print(c)

例中数组shape为 3×2×4
输出结果为：
0 对应shape 2×4
1 对应shape 3×4
2 对应shape 3×2

[[2 1 0 2]                                                                                                               
 [2 1 1 2]]

[[0 0 0 1]                                                                                                               
 [0 0 1 1]                                                                                                               
 [0 1 0 0]]

[[3 3]                                                                                                                   
 [1 2]                                                                                                                   
 [3 3]]  
 

https://blog.csdn.net/weixin_42755982/article/details/104542538