Pytorch - 在 softmax 层之后选择最佳概率

Question

我一直在使用Pytorch 0.4.0逻辑回归模型，在我的输入为高维和我的输出必须是一个标量- 0，1或2。

我使用一个线性层与一个 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]])