Numpy 从softmax输出到类预测

从Softmax输出到类预测有没有简单的方法

比如说,

由此：

[0.83128697, 0.06161868, 0.10709436]

为此：

[1, 0, 0]

---

您可以使用

np.argmax

检索最大值的索引：

import numpy as np
a = [0.83128697, 0.06161868, 0.10709436]

r = np.zeros(len(a)) # a.size if a is a numpy array
r[np.argmax(a)]=1

r
array([1., 0., 0.])