np.argmax()函数用法，返回最大值的索引

这篇文章写的太好了

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

import numpy as np

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]]]

three_dim_array = np.array(three_dim_array)
print(three_dim_array.shape)

a = np.argmax(three_dim_array, axis = 0)
print(a)
b = np.argmax(three_dim_array, axis = 1)
print(b)
c = np.argmax(three_dim_array, axis = 2)
print(c)

输出
(3, 2, 4)
[[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]]