Pytorch - выбирает наилучшую вероятность после слоя softmax

Question

У меня есть модель логистической регрессии, использующая Pytorch 0.4.0, где мой ввод является многомерным, а мой вывод должен быть скалярным - 0, 1 или 2.

I 'm используя линейный слой в сочетании со слоем softmax для возврата тензора n x 3, где каждый столбец представляет вероятность попадания входных данных в один из трех классов (0, 1 или 2).

Однако я должен вернуть тензор n x 1, поэтому мне нужно каким-то образом выбрать наибольшую вероятность для каждого входа и создать тензор, указывающий, какой класс имел наибольшую вероятность.Как я могу добиться этого с помощью Pytorch?

Чтобы проиллюстрировать это, мой Softmax выводит следующее:

[[0.2, 0.1, 0.7],
 [0.6, 0.2, 0.2],
 [0.1, 0.8, 0.1]]

И я должен вернуть это:

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

Answer 1

torch.argmax(), вероятно, то, что вы хотите:

import torch

x = torch.FloatTensor([[0.2, 0.1, 0.7],
                       [0.6, 0.2, 0.2],
                       [0.1, 0.8, 0.1]])

y = torch.argmax(x, dim=1)
print(y.detach())
# tensor([ 2,  0,  1])

# If you want to reshape:
y = y.view(1, -1)
print(y.detach())
# tensor([[ 2,  0,  1]])