Как я могу использовать np argmax в массиве списков?

Question

У меня есть данные y_hat, которые выглядят так:

[[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.

Answer 1

Вот один путь после argmax с numpy трансляцией

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

Или

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

Answer 2

np.argmax принимает аргумент ключевого слова axis .Используйте это.

Это axis=0 для столбцов, axis=1 для строк.