Python 如何在列表数组中使用np argmax？

我有一些数据看起来像：

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

我想得到每一行的

argmax

，这样我就得到了如下向量：

[[3]
 [8]
 [8]
 ...
 [5]
 [1]
 [7]]

如果我只执行

np.argmax（y_hat）

，它将返回

1

。用那个

---

列的轴是

axis=0

，行的轴是

axis=1

。

下面是

argmax

之后的一种方式，使用

numpy

广播

a.argmax(axis = 1)[:,None]

或

---

通常我们将

y_hat

描述为一个2d数组，如果您想在“行”之间求和，他可能还想使用

keepdims

参数。

a[:,None].argmax(-1)